Belle II Software  release-08-01-10
savingFlippingVariables.py
1 
8 
9 import ROOT
10 from ROOT import Belle2
11 import math
12 import tracking.harvest.harvesting as harvesting
13 import tracking.harvest.refiners as refiners
14 import tracking.harvest.peelers as peelers
15 from tracking.validation.utilities import getObjectList
16 ROOT.gSystem.Load("libtracking")
17 
18 
19 class 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 
27  self.checkObjcheckObj = checkObj
28 
29 
30  self.mcRecoTracksmcRecoTracks = "MCRecoTracks"
31 
32 
33  self.track_match_look_uptrack_match_look_up = None
34 
35 
36  self.outputnameoutputname = 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_uptrack_match_look_up = Belle2.TrackMatchLookUp(self.mcRecoTracksmcRecoTracks, self.checkObjcheckObj)
42  output_tfile = ROOT.TFile(self.outputnameoutputname, "RECREATE")
43  self.outputnameoutputname = 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.checkObjcheckObj)
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_uptrack_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 
268  save_tree = refiners.save_tree(name="data")
A (simplified) python wrapper for StoreArray.
Definition: PyStoreArray.h:72
Class to provide convenient methods to look up matching information between pattern recognition and M...
track_match_look_up
Reference to the track match lookup object reading the relation information.
def __init__(self, name, contact=None, checkObj='RecoTracks', output_file_name='flip-refit-MVA1.root')
mcRecoTracks
Name of the StoreArray of the mc tracks.