Belle II Software development
savingFlippingVariablesFor2ndMVA.py
1
8
9import ROOT
10from ROOT import Belle2
11import math
12import numpy as np
13import tracking.harvest.harvesting as harvesting
14import tracking.harvest.refiners as refiners
15ROOT.gSystem.Load("libtracking")
16
17
18class Saving2ndMVAData(harvesting.HarvestingModule):
19 """ A dedicated module to save the variables using in flipping steps"""
20
21 def __init__(self, name, contact=None, checkObj='RecoTracks', output_file_name='flip-refit-MVA2.root'):
22 """Constructor"""
23 super().__init__(foreach=checkObj, name=name, contact=contact, output_file_name=output_file_name)
24
25
26 self.checkObj = checkObj
27
28
29 self.mcRecoTracks = "MCRecoTracks"
30
31
32 self.track_match_look_up = None
33
34
35 self.outputname = output_file_name
36
37 def initialize(self):
38 """Initialization at the start of the event processing"""
39 super().initialize()
40 self.track_match_look_up = Belle2.TrackMatchLookUp(self.mcRecoTracks, self.checkObj)
41 output_tfile = ROOT.TFile(self.outputname, "RECREATE")
42 self.outputname = output_tfile
43
44 def prepare(self):
45 """preparation at the start of each event.
46 make sure the checkObj exist
47 """
48 super().prepare()
49 checkDatas = Belle2.PyStoreArray(self.checkObj)
50 if (not checkDatas):
51 return False
52
53 def pick(self, recoTrack):
54 """pick every recoTrack"""
55 return True
56
57 def peel(self, recoTrack):
58 """store the information for each recoTrack"""
59 track_match_look_up = self.track_match_look_up
60 nan = float('nan')
61 flipped_pz_estimate = nan
62 y_variance = nan
63 tan_lambda_estimate = nan
64 d0_variance = nan
65 x_variance = nan
66 z_estimate = nan
67 phi0_variance = nan
68 px_variance = nan
69 pz_estimate = nan
70 p_value = nan
71 pt_estimate = nan
72 y_estimate = nan
73 d0_estimate = nan
74 x_estimate = nan
75 py_variance = nan
76 pz_variance = nan
77 omega_variance = nan
78 tan_lambda_variance = nan
79 z_variance = nan
80 omega_estimate = nan
81 pt_resolution = nan
82 px_estimate = nan
83 pt_variance = nan
84 phi0_estimate = nan
85 flipped_z_estimate = nan
86 py_estimate = nan
87 flipped_z_variance = nan
88 flipped_pz_variance = nan
89 flipped_pt_variance = nan
90 flipped_py_estimate = nan
91 z0_variance = nan
92 flipped_p_value = nan
93 flipped_px_variance = nan
94 flipped_py_variance = nan
95 flipped_x_estimate = nan
96 quality_flip_indicator = nan
97 quality_2ndflip_indicator = nan
98 isPrimary_misID = False
99 ismatched = False
100 ismatched_CC = False
101 ismatched_WC = False
102 isclone_CC = False
103 isclone_WC = False
104 isclone = False
105 isbackground = False
106 isghost = False
107 isprimary = False
108 charge_truth = nan
109 track_charge = nan
110 inGoingArmTime = nan
111 inGoingArmTimeError = nan
112 outGoingArmTime = nan
113 outGoingArmTimeError = nan
114 InOutArmTimeDifference = nan
115 InOutArmTimeDifferenceError = nan
116
117 if (recoTrack):
118 mc_particle = track_match_look_up.getRelatedMCParticle(recoTrack)
119 fit_result = track_match_look_up.getRelatedTrackFitResult(recoTrack)
120
121 inGoingArmTime = recoTrack.getIngoingArmTime()
122 inGoingArmTimeError = recoTrack.getIngoingArmTimeError()
123 outGoingArmTime = recoTrack.getOutgoingArmTime()
124 outGoingArmTimeError = recoTrack.getOutgoingArmTimeError()
125 InOutArmTimeDifference = recoTrack.getInOutArmTimeDifference()
126 InOutArmTimeDifferenceError = recoTrack.getInOutArmTimeDifferenceError()
127
128 quality_flip_indicator = recoTrack.getFlipQualityIndicator()
129 quality_2ndflip_indicator = recoTrack.get2ndFlipQualityIndicator()
130
131 ismatched = track_match_look_up.isAnyChargeMatchedPRRecoTrack(recoTrack)
132 ismatched_CC = track_match_look_up.isCorrectChargeMatchedPRRecoTrack(recoTrack)
133 ismatched_WC = track_match_look_up.isWrongChargeMatchedPRRecoTrack(recoTrack)
134
135 isclone = track_match_look_up.isAnyChargeClonePRRecoTrack(recoTrack)
136 isclone_CC = track_match_look_up.isCorrectChargeClonePRRecoTrack(recoTrack)
137 isclone_WC = track_match_look_up.isWrongChargeClonePRRecoTrack(recoTrack)
138
139 isbackground = track_match_look_up.isBackgroundPRRecoTrack(recoTrack)
140 isghost = track_match_look_up.isGhostPRRecoTrack(recoTrack)
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 charge_truth != track_charge:
148 isPrimary_misID = True
149
150 recoTrack_flipped = recoTrack.getRelated("RecoTracks_flipped")
151 if recoTrack_flipped:
152 track_flipped = recoTrack_flipped.getRelated("Tracks_flipped")
153 if track_flipped:
154 fit_result_flipped = track_flipped.getTrackFitResultWithClosestMassByName(
155 Belle2.Const.pion, "TrackFitResults_flipped")
156 if (fit_result and fit_result_flipped):
157 cov6 = fit_result.getCovariance6()
158 mom = fit_result.getMomentum()
159 pos = fit_result.getPosition()
160
161 pt_estimate = mom.Rho()
162 pt_variance = np.divide(mom.X() ** 2 * cov6(3, 3) + mom.Y() ** 2 * cov6(4, 4) -
163 2 * mom.X() * mom.Y() * cov6(3, 4), mom.Perp2())
164 pt_resolution = np.divide(pt_variance, pt_estimate)
165
166 tan_lambda_estimate = fit_result.getCotTheta()
167 omega_estimate = fit_result.getOmega()
168 phi0_estimate = fit_result.getPhi() % (2.0 * math.pi)
169 d0_estimate = fit_result.getD0()
170 d0_variance = fit_result.getCov()[0]
171 z0_variance = fit_result.getCov()[12]
172 phi0_variance = fit_result.getCov()[5]
173 omega_variance = fit_result.getCov()[9]
174 tan_lambda_variance = fit_result.getCov()[14]
175 x_estimate = pos.X()
176 y_estimate = pos.Y()
177 z_estimate = pos.Z()
178 x_variance = cov6(0, 0)
179 y_variance = cov6(1, 1)
180 z_variance = cov6(2, 2)
181 px_estimate = mom.X()
182 py_estimate = mom.Y()
183 pz_estimate = cov6(5, 5)
184 px_variance = cov6(3, 3)
185 py_variance = cov6(4, 4)
186 pz_variance = cov6(5, 5)
187 p_value = fit_result.getPValue()
188
189 cov6_flipped = fit_result_flipped.getCovariance6()
190 mom_flipped = fit_result_flipped.getMomentum()
191 pos_flipped = fit_result_flipped.getPosition()
192
193 flipped_pt_variance = np.divide(
194 mom_flipped.X() ** 2 * cov6_flipped(3, 3) +
195 mom_flipped.Y() ** 2 * cov6_flipped(4, 4) -
196 2 * mom_flipped.X() * mom_flipped.Y() * cov6_flipped(3, 4),
197 mom_flipped.Perp2())
198
199 flipped_z_estimate = pos_flipped.Z()
200 flipped_pz_estimate = mom_flipped.Z()
201 flipped_py_estimate = mom_flipped.Y()
202
203 flipped_x_estimate = pos_flipped.X()
204 flipped_z_variance = cov6_flipped(2, 2)
205 flipped_pz_variance = cov6_flipped(5, 5)
206 flipped_px_variance = cov6_flipped(3, 3)
207 flipped_py_variance = cov6_flipped(4, 4)
208 flipped_p_value = fit_result_flipped.getPValue()
209
210 crops = dict(
211 flipped_pz_estimate=flipped_pz_estimate,
212 y_variance=y_variance,
213 tan_lambda_estimate=tan_lambda_estimate,
214 d0_variance=d0_variance,
215 x_variance=x_variance,
216 z_estimate=z_estimate,
217 phi0_variance=phi0_variance,
218 px_variance=px_variance,
219 pz_estimate=pz_estimate,
220 p_value=p_value,
221 pt_estimate=pt_estimate,
222 y_estimate=y_estimate,
223 d0_estimate=d0_estimate,
224 x_estimate=x_estimate,
225 py_variance=py_variance,
226 pz_variance=pz_variance,
227 omega_variance=omega_variance,
228 tan_lambda_variance=tan_lambda_variance,
229 z_variance=z_variance,
230 omega_estimate=omega_estimate,
231 pt_resolution=pt_resolution,
232 px_estimate=px_estimate,
233 pt_variance=pt_variance,
234 phi0_estimate=phi0_estimate,
235 flipped_z_estimate=flipped_z_estimate,
236 py_estimate=py_estimate,
237 flipped_z_variance=flipped_z_variance,
238 flipped_pz_variance=flipped_pz_variance,
239 flipped_pt_variance=flipped_pt_variance,
240 flipped_py_estimate=flipped_py_estimate,
241 z0_variance=z0_variance,
242 flipped_p_value=flipped_p_value,
243 flipped_px_variance=flipped_px_variance,
244 flipped_py_variance=flipped_py_variance,
245 flipped_x_estimate=flipped_x_estimate,
246 quality_flip_indicator=quality_flip_indicator,
247 quality_2ndflip_indicator=quality_2ndflip_indicator,
248 isPrimary_misID=isPrimary_misID,
249 ismatched=ismatched,
250 ismatched_CC=ismatched_CC,
251 ismatched_WC=ismatched_WC,
252 isclone_CC=isclone_CC,
253 isclone_WC=isclone_WC,
254 isclone=isclone,
255 isbackground=isbackground,
256 isghost=isghost,
257 isprimary=isprimary,
258 charge_truth=charge_truth,
259 track_charge=track_charge,
260 inGoingArmTime=inGoingArmTime,
261 inGoingArmTimeError=inGoingArmTimeError,
262 outGoingArmTime=outGoingArmTime,
263 outGoingArmTimeError=outGoingArmTimeError,
264 InOutArmTimeDifference=InOutArmTimeDifference,
265 InOutArmTimeDifferenceError=InOutArmTimeDifferenceError,
266 )
267 return crops
268
269
270 save_tree = refiners.save_tree(name="data")
Belle2::PyStoreArray
A (simplified) python wrapper for StoreArray.
Definition: PyStoreArray.h:72
Belle2::TrackMatchLookUp
Class to provide convenient methods to look up matching information between pattern recognition and M...
Definition: TrackMatchLookUp.h:33
savingFlippingVariablesFor2ndMVA.Saving2ndMVAData.track_match_look_up
track_match_look_up
Reference to the track match lookup object reading the relation information.
Definition: savingFlippingVariablesFor2ndMVA.py:32
savingFlippingVariablesFor2ndMVA.Saving2ndMVAData.__init__
def __init__(self, name, contact=None, checkObj='RecoTracks', output_file_name='flip-refit-MVA2.root')
Definition: savingFlippingVariablesFor2ndMVA.py:21