Belle II Software development
savingFlippingVariables.py
1
8
9import ROOT
10from ROOT import Belle2
11import math
12import tracking.harvest.harvesting as harvesting
13import tracking.harvest.refiners as refiners
14import tracking.harvest.peelers as peelers
15from tracking.validation.utilities import getObjectList
16ROOT.gSystem.Load("libtracking")
17
18
19class Saving1stMVAData(harvesting.HarvestingModule):
20 """ A dedicated module to save the variables using in flipping steps"""
21
22 def __init__(self, name, contact=None, checkObj='RecoTracks', output_file_name='flip-refit-MVA1.root'):
23 """Constructor"""
24 super().__init__(foreach=checkObj, name=name, contact=contact, output_file_name=output_file_name)
25
26 ## Name of the Obj to be picked
27 self.checkObj = checkObj
28
29 ## Name of the StoreArray of the mc tracks
30 self.mcRecoTracks = "MCRecoTracks"
31
32 ## Reference to the track match lookup object reading the relation information
33 self.track_match_look_up = None
34
35 ## Name of the output file
36 self.outputname = output_file_name
37
38 def initialize(self):
39 """Initialization at the start of the event processing"""
40 super().initialize()
41 self.track_match_look_up = Belle2.TrackMatchLookUp(self.mcRecoTracks, self.checkObj)
42 output_tfile = ROOT.TFile(self.outputname, "RECREATE")
43 self.outputname = output_tfile
44
45 def prepare(self):
46 """preparation at the start of each event.
47 make sure the checkObj exist
48 """
49 super().prepare()
50 checkDatas = Belle2.PyStoreArray(self.checkObj)
51 if (not checkDatas):
52 return False
53
54 def pick(self, recoTrack):
55 """pick every recoTrack"""
56 return True
57
58 def peel(self, recoTrack):
59 """store the information for each recoTrack"""
60 track_match_look_up = self.track_match_look_up
61 nan = float('nan')
62 d0_variance = nan
63 seed_pz_estimate = nan
64 n_hits = nan
65 z0_estimate = nan
66 seed_pz_variance = nan
67 phi0_variance = nan
68 seed_z_estimate = nan
69 tan_lambda_estimate = nan
70 omega_variance = nan
71 seed_tan_lambda_estimate = nan
72 d0_estimate = nan
73 seed_pt_estimate = nan
74 cdc_qualityindicator = nan
75 omega_estimate = nan
76 z0_variance = nan
77 seed_x_estimate = nan
78 seed_y_estimate = nan
79 seed_pt_resolution = nan
80 seed_py_variance = nan
81 seed_d0_estimate = nan
82 seed_omega_variance = nan
83 tan_lambda_variance = nan
84 svd_layer6_clsTime = nan
85 seed_tan_lambda_variance = nan
86 seed_z_variance = nan
87 n_svd_hits = nan
88 phi0_estimate = nan
89 n_cdc_hits = nan
90 n_pxd_hits = nan
91 svd_layer3_positionSigma = nan
92 first_cdc_layer = nan
93 last_cdc_layer = nan
94 ndf_hits = nan
95 isPrimary_misID = False
96 ismatched = False
97 ismatched_CC = False
98 ismatched_WC = False
99 isclone_CC = False
100 isclone_WC = False
101 isclone = False
102 isbackground = False
103 isghost = False
104 isprimary = False
105 charge_truth = nan
106 inGoingArmTime = nan
107 inGoingArmTimeError = nan
108 outGoingArmTime = nan
109 outGoingArmTimeError = nan
110 InOutArmTimeDifference = nan
111 InOutArmTimeDifferenceError = nan
112 pt_estimate = nan
113 track_charge = nan
114 quality_flip_indicator = nan
115 quality_2ndflip_indicator = nan
116
117 if (recoTrack):
118
119 mc_particle = track_match_look_up.getRelatedMCParticle(recoTrack)
120 fit_result = track_match_look_up.getRelatedTrackFitResult(recoTrack)
121
122 inGoingArmTime = recoTrack.getIngoingArmTime()
123 inGoingArmTimeError = recoTrack.getIngoingArmTimeError()
124 outGoingArmTime = recoTrack.getOutgoingArmTime()
125 outGoingArmTimeError = recoTrack.getOutgoingArmTimeError()
126 InOutArmTimeDifference = recoTrack.getInOutArmTimeDifference()
127 InOutArmTimeDifferenceError = recoTrack.getInOutArmTimeDifferenceError()
128
129 ismatched = track_match_look_up.isAnyChargeMatchedPRRecoTrack(recoTrack)
130 ismatched_CC = track_match_look_up.isCorrectChargeMatchedPRRecoTrack(recoTrack)
131 ismatched_WC = track_match_look_up.isWrongChargeMatchedPRRecoTrack(recoTrack)
132
133 isclone = track_match_look_up.isAnyChargeClonePRRecoTrack(recoTrack)
134 isclone_CC = track_match_look_up.isCorrectChargeClonePRRecoTrack(recoTrack)
135 isclone_WC = track_match_look_up.isWrongChargeClonePRRecoTrack(recoTrack)
136
137 isbackground = track_match_look_up.isBackgroundPRRecoTrack(recoTrack)
138 isghost = track_match_look_up.isGhostPRRecoTrack(recoTrack)
139 quality_flip_indicator = recoTrack.getFlipQualityIndicator()
140 quality_2ndflip_indicator = recoTrack.get2ndFlipQualityIndicator()
141
142 if mc_particle and fit_result:
143 isprimary = bool(mc_particle.hasStatus(Belle2.MCParticle.c_PrimaryParticle))
144 charge_truth = mc_particle.getCharge()
145 track_charge = fit_result.getChargeSign()
146 if isprimary:
147 if mc_particle.getCharge() != track_charge:
148 isPrimary_misID = True
149
150 cdc_track_cand = recoTrack.getRelated('CDCRecoTracks')
151 if cdc_track_cand:
152 cdc_qualityindicator = cdc_track_cand.getQualityIndicator()
153
154 if fit_result:
155 omega_estimate = fit_result.getOmega()
156 z0_estimate = fit_result.getZ0()
157 d0_estimate = fit_result.getD0()
158 phi0_estimate = fit_result.getPhi() % (2.0 * math.pi)
159 tan_lambda_estimate = fit_result.getCotTheta()
160
161 mom = fit_result.getMomentum()
162 pt_estimate = mom.Rho()
163
164 d0_variance = fit_result.getCov()[0]
165 z0_variance = fit_result.getCov()[12]
166 phi0_variance = fit_result.getCov()[5]
167 omega_variance = fit_result.getCov()[9]
168
169 reco_svdcdc_track = recoTrack.getRelated("SVDCDCRecoTracks")
170
171 seed_fit_result = peelers.get_seed_track_fit_result(reco_svdcdc_track)
172 seed_mom = seed_fit_result.getMomentum()
173 seed_pos = seed_fit_result.getPosition()
174 seed_cov6 = seed_fit_result.getCovariance6()
175 seed_tan_lambda_estimate = seed_fit_result.getCotTheta()
176
177 seed_pz_estimate = seed_mom.Z()
178 seed_pz_variance = seed_cov6(5, 5)
179 seed_z_estimate = seed_pos.Z()
180 seed_x_estimate = seed_pos.X()
181 seed_y_estimate = seed_pos.Y()
182
183 seed_pt_estimate = seed_mom.Rho()
184 seed_py_variance = seed_cov6(4, 4)
185 seed_d0_estimate = seed_fit_result.getD0()
186 seed_omega_variance = seed_fit_result.getCov()[9]
187 seed_tan_lambda_variance = seed_fit_result.getCov()[14]
188 seed_z_variance = seed_cov6(2, 2)
189
190 tan_lambda_variance = seed_fit_result.getCov()[14]
191 for svd_hit in getObjectList(recoTrack.getSVDHitList()):
192 if svd_hit.getSensorID().getLayerNumber() == 3:
193 svd_layer3_positionSigma = svd_hit.getPositionSigma()
194 if svd_hit.getSensorID().getLayerNumber() == 6:
195 svd_layer6_clsTime = svd_hit.getClsTime()
196
197 cdc_hits = [hit.getICLayer() for hit in getObjectList(recoTrack.getCDCHitList())]
198 if cdc_hits:
199 first_cdc_layer = min(cdc_hits)
200 last_cdc_layer = max(cdc_hits)
201
202 n_cdc_hits = recoTrack.getNumberOfCDCHits()
203 n_svd_hits = recoTrack.getNumberOfSVDHits()
204 n_pxd_hits = recoTrack.getNumberOfPXDHits()
205
206 n_hits = n_pxd_hits + n_svd_hits + n_cdc_hits
207 ndf_hits = 2 * n_pxd_hits + n_svd_hits + n_cdc_hits
208
209 crops = dict(
210 d0_variance=d0_variance,
211 seed_pz_estimate=seed_pz_estimate,
212 n_hits=n_hits,
213 z0_estimate=z0_estimate,
214 seed_pz_variance=seed_pz_variance,
215 phi0_variance=phi0_variance,
216 seed_z_estimate=seed_z_estimate,
217 tan_lambda_estimate=tan_lambda_estimate,
218 omega_variance=omega_variance,
219 seed_tan_lambda_estimate=seed_tan_lambda_estimate,
220 d0_estimate=d0_estimate,
221 seed_pt_estimate=seed_pt_estimate,
222 cdc_qualityindicator=cdc_qualityindicator,
223 omega_estimate=omega_estimate,
224 z0_variance=z0_variance,
225 seed_x_estimate=seed_x_estimate,
226 seed_y_estimate=seed_y_estimate,
227 seed_pt_resolution=seed_pt_resolution,
228 seed_py_variance=seed_py_variance,
229 seed_d0_estimate=seed_d0_estimate,
230 seed_omega_variance=seed_omega_variance,
231 tan_lambda_variance=tan_lambda_variance,
232 svd_layer6_clsTime=svd_layer6_clsTime,
233 seed_tan_lambda_variance=seed_tan_lambda_variance,
234 seed_z_variance=seed_z_variance,
235 n_svd_hits=n_svd_hits,
236 phi0_estimate=phi0_estimate,
237 n_cdc_hits=n_cdc_hits,
238 n_pxd_hits=n_pxd_hits,
239 svd_layer3_positionSigma=svd_layer3_positionSigma,
240 first_cdc_layer=first_cdc_layer,
241 last_cdc_layer=last_cdc_layer,
242 ndf_hits=ndf_hits,
243 isPrimary_misID=isPrimary_misID,
244 ismatched=ismatched,
245 ismatched_CC=ismatched_CC,
246 ismatched_WC=ismatched_WC,
247 isclone_CC=isclone_CC,
248 isclone_WC=isclone_WC,
249 isclone=isclone,
250 isbackground=isbackground,
251 isghost=isghost,
252 isprimary=isprimary,
253 charge_truth=charge_truth,
254 track_charge=track_charge,
255 inGoingArmTime=inGoingArmTime,
256 inGoingArmTimeError=inGoingArmTimeError,
257 outGoingArmTime=outGoingArmTime,
258 outGoingArmTimeError=outGoingArmTimeError,
259 InOutArmTimeDifference=InOutArmTimeDifference,
260 InOutArmTimeDifferenceError=InOutArmTimeDifferenceError,
261 pt_estimate=pt_estimate,
262 quality_flip_indicator=quality_flip_indicator,
263 quality_2ndflip_indicator=quality_2ndflip_indicator,
264 )
265 return crops
266
267 ## save a tree of all the collected variables
268 save_tree = refiners.save_tree(name="data")
savingFlippingVariables.Saving1stMVAData.__init__
__init__(self, name, contact=None, checkObj='RecoTracks', output_file_name='flip-refit-MVA1.root')
Definition savingFlippingVariables.py:22