Belle II Software  release-08-01-10
savingFlippingVariablesFor2ndMVA.py
1 
8 
9 import ROOT
10 from ROOT import Belle2
11 import math
12 import numpy as np
13 import tracking.harvest.harvesting as harvesting
14 import tracking.harvest.refiners as refiners
15 ROOT.gSystem.Load("libtracking")
16 
17 
18 class 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.checkObjcheckObj = checkObj
27 
28 
29  self.mcRecoTracksmcRecoTracks = "MCRecoTracks"
30 
31 
32  self.track_match_look_uptrack_match_look_up = None
33 
34 
35  self.outputnameoutputname = 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_uptrack_match_look_up = Belle2.TrackMatchLookUp(self.mcRecoTracksmcRecoTracks, self.checkObjcheckObj)
41  output_tfile = ROOT.TFile(self.outputnameoutputname, "RECREATE")
42  self.outputnameoutputname = 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.checkObjcheckObj)
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_uptrack_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