Belle II Software development
collections.py
1
8
9import basf2
10import rawdata as raw
11import reconstruction as reco
12import modularAnalysis as ana
13import vertex as vtx
14
15from collections import namedtuple
16
17MillepedeCollection = namedtuple('MillepedeCollection', ['name', 'files', 'path', 'params'])
18
19
20def make_collection(name, files=None, path=None, **argk):
21 """
22 Handy function to make a collection configuration
23 to be passed in 'collections' argument of the create(...) function
24
25 Parameters
26 ----------
27 name : str
28 Collection name
29 files : list(str)
30 List of input data files
31 path : basf2.Path
32 pre-collector path
33 argk : dict
34 Dictionary of collector parameters specific for collection
35
36 Returns
37 -------
38 namedtuple('MillepedeCollection', ['name', 'files', 'path', 'params'])
39
40 """
41 if files is None:
42 files = []
43 if path is None:
44 path = basf2.Path()
45
46 return MillepedeCollection(name=name, path=path, files=files, params=argk)
47
48
49def physicsTracks(name="physicsTracks", files=None, add_unpackers=True, klm=False, prescale=1.):
50 """
51 Standard collection of all RecoTracks with standard reconstruction
52
53 Parameters
54 ----------
55 name : str
56 Collection name
57 files : list(str)
58 List of input data files
59 add_unpackers : bool
60 Whether to add unpacking (set to False for MC)
61 klm : bool
62 Whether to add muid hits to the track fit
63 prescale : float
64 Process only 'prescale' fraction of events
65 """
66 path = basf2.create_path()
67
68 path.add_module('Progress')
69 path.add_module('RootInput')
70 if prescale != 1.:
71 path.add_module('Prescale', prescale=prescale).if_false(basf2.Path(), basf2.AfterConditionPath.END)
72 path.add_module('Gearbox')
73 path.add_module('Geometry')
74
75 if add_unpackers:
76 raw.add_unpackers(path)
77
78 reco.add_reconstruction(path, pruneTracks=False, add_muid_hits=klm)
79
80 # teporary (?) fix : TODO: still needed?
81 tmp = basf2.create_path()
82 for m in path.modules():
83 if m.name() == "PXDPostErrorChecker":
84 m.param('CriticalErrorMask', 0)
85 if m.name() == "SVDSpacePointCreator":
86 m.param("MinClusterTime", -999)
87 tmp.add_module(m)
88 path = tmp
89 path.add_module('DAFRecoFitter')
90
91 return make_collection(name, files=files, path=path, tracks=['RecoTracks'])
92
93
94def cosmicTracks(name="cosmicTracks",
95 files=None,
96 add_unpackers=True,
97 skim_hlt_cosmic=False,
98 cut='[z0 <= 57. or abs(d0) >= 26.5] and abs(dz) > 0.4 and nTracks == 1',
99 klm=False,
100 prescale=1.):
101 """
102 Standard collection of all RecoTracks with cosmic reconstruction
103
104 Parameters
105 ----------
106 name : str
107 Collection name
108 files : list(str)
109 List of input data files
110 add_unpackers : bool
111 Whether to add unpacking (set to False for MC)
112 skim_hlt_cosmic : bool
113 Whether to add TriggerSkim module and process only events with cosmic TRG
114 cut : str
115 Cut string to select GOOD events. By default set to avoid region of poorly described magnetic
116 field around QCS magnets + remove the 'background' from physics around IP
117 klm : bool
118 Whether to add muid hits to the track fit
119 prescale : float
120 Process only 'prescale' fraction of events
121 """
122 path = basf2.create_path()
123
124 path.add_module('Progress')
125 path.add_module('RootInput')
126 if prescale != 1.:
127 path.add_module('Prescale', prescale=prescale).if_false(basf2.Path(), basf2.AfterConditionPath.END)
128
129 if skim_hlt_cosmic:
130 path.add_module(
131 "TriggerSkim",
132 triggerLines=["software_trigger_cut&filter&cosmic"]).if_value(
133 "==0",
134 basf2.Path(),
135 basf2.AfterConditionPath.END)
136
137 path.add_module('Gearbox')
138 path.add_module('Geometry')
139
140 if add_unpackers:
141 raw.add_unpackers(path)
142
143 path.add_module('SetupGenfitExtrapolation')
144 reco.add_cosmics_reconstruction(
145 path,
146 pruneTracks=False,
147 skipGeometryAdding=True,
148 merge_tracks=True,
149 add_muid_hits=klm
150 )
151
152 # teporary (?) fix : TODO: still needed?
153 tmp = basf2.create_path()
154 for m in path.modules():
155 if m.name() == "PXDPostErrorChecker":
156 m.param('CriticalErrorMask', 0)
157 if m.name() == "SVDSpacePointCreator":
158 m.param("MinClusterTime", -999)
159 tmp.add_module(m)
160 path = tmp
161
162 path.add_module('SetRecoTrackMomentum', automatic=True)
163 path.add_module('DAFRecoFitter', resortHits=False, pdgCodesToUseForFitting=[13])
164
165 if cut is not None:
166 ana.fillParticleList('mu+:good_cosmics', cut, path=path)
167 # ana.variablesToNtuple('mu+:good_cosmics', variables=track_variables, filename='analysis_cosmics.root', path=path)
168 path.add_module('SkimFilter', particleLists=['mu+:good_cosmics']).if_false(basf2.create_path())
169
170 return make_collection(name, files=files, path=path, tracks=['RecoTracks'])
171
172
173def diMuonsIP(
174 name="diMuonsIP",
175 files=None,
176 add_unpackers=True,
177 skim_mumu_2trk=False,
178 muon_cut='p > 1.0 and abs(dz) < 2.0 and dr < 0.5',
179 dimuon_cut='',
180 klm=False,
181 prescale=1.):
182 """
183 Di-muons with vertex+beam constraint collection
184
185 Parameters
186 ----------
187 name : str
188 Collection name
189 files : list(str)
190 List of input data files
191 add_unpackers : bool
192 Whether to add unpacking (set to False for MC)
193 skim_hlt_cosmic : bool
194 Whether to add TriggerSkim module and process only events with accept_mumu_2trk TRG
195 muon_cut : str
196 Cut string to select daughter muons
197 dimuon_cut : str
198 Cut string to apply for reconstructed di-muon decay
199 klm : bool
200 Whether to add muid hits to the track fit
201 prescale : float
202 Process only 'prescale' fraction of events
203 """
204 path = basf2.create_path()
205 path.add_module('Progress')
206 path.add_module('RootInput')
207 if prescale != 1.:
208 path.add_module('Prescale', prescale=prescale).if_false(basf2.Path(), basf2.AfterConditionPath.END)
209
210 path.add_module('Gearbox')
211 path.add_module('Geometry')
212
213 if skim_mumu_2trk:
214 path.add_module(
215 "TriggerSkim",
216 triggerLines=["software_trigger_cut&skim&accept_mumu_2trk"]).if_value(
217 "==0",
218 basf2.Path(),
219 basf2.AfterConditionPath.END)
220
221 if add_unpackers:
222 raw.add_unpackers(path)
223
224 reco.add_reconstruction(path, pruneTracks=False, add_muid_hits=klm)
225
226 # teporary (?) fix : TODO: still needed?
227 tmp = basf2.create_path()
228 for m in path.modules():
229 if m.name() == "PXDPostErrorChecker":
230 m.param('CriticalErrorMask', 0)
231 if m.name() == "SVDSpacePointCreator":
232 m.param("MinClusterTime", -999)
233 tmp.add_module(m)
234 path = tmp
235
236 path.add_module('DAFRecoFitter', pdgCodesToUseForFitting=[13])
237
238 ana.fillParticleList(f"mu+:{name}", muon_cut, path=path)
239 ana.reconstructDecay(f"Upsilon(4S):{name} -> mu+:{name} mu-:{name}", dimuon_cut, path=path)
240
241 vtx.raveFit(f"Upsilon(4S):{name}", 0.001, daughtersUpdate=True, silence_warning=True, path=path)
242
243 return make_collection(name, files=files, path=path, primaryVertices=[f"Upsilon(4S):{name}"])