Belle II Software  release-05-01-25
asymmetriesInVariablesPlots.py
1 #!/usr/bin/env python3
2 # -*- coding: utf-8 -*-
3 
4 
18 
19 import ROOT
20 from ROOT import Belle2
21 import sys
22 import flavorTagger as ft
23 from defaultEvaluationParameters import r_subsample, r_size, categories
24 from inputVariablesPlots import variablesPlotParamsDict
25 import basf2_mva
26 from array import array
27 
28 import numpy as np
29 import matplotlib as mpl
30 mpl.use('Agg')
31 mpl.rcParams.update({'font.size': 22})
32 mpl.rcParams['text.usetex'] = True
33 mpl.rcParams['text.latex.preamble'] = [r"\usepackage{amsmath}"]
34 import matplotlib.pyplot as plt
35 from matplotlib.ticker import FormatStrFormatter
36 import math
37 import glob
38 import os
39 
40 
41 if len(sys.argv) != 5:
42  sys.exit("Must provide 4 arguments: [Belle or Belle2] [BGx0 or BGx1] [training] [workingDirectory]"
43  )
44 
45 belleOrBelle2 = sys.argv[1]
46 MCtype = str(sys.argv[2])
47 decayChannelTrainedOn = str(sys.argv[3]) # Decay channel of the weight files "JPsiKs" or "nunubar"
48 workingDirectory = sys.argv[4]
49 
50 filesDirectory = workingDirectory + '/FlavorTagging/TrainedMethods'
51 
52 if decayChannelTrainedOn == 'JPsiKs':
53  decayChannelTrainedOn = 'JpsiKs_mu'
54 
55 weightFiles = 'B2' + decayChannelTrainedOn + MCtype
56 
57 ROOT.TH1.SetDefaultSumw2()
58 
59 allInputVariables = []
60 
61 belleOrBelle2Flag = belleOrBelle2
62 
63 identifiersExtraInfosDict = dict()
64 identifiersExtraInfosKaonPion = []
65 
66 dBw = 50
67 
68 pBins = 50
69 fBins = 100
70 
71 unitImp = "mm"
72 
73 if belleOrBelle2 == "Belle":
74  unitImp = "cm"
75 
76 
77 # particle Label, PDG for negative, PDG for positive particle
78 
79 particleConditions = {'e+:inRoe': ['e', " mcPDG > 0 ", " mcPDG < 0 "], 'mu+:inRoe': [r'\mu', " mcPDG > 0 ", " mcPDG < 0 "],
80  'lepton+:inRoe': ['l', " mcPDG > 0 ", " mcPDG < 0 "],
81  'K+:inRoe': ['K', " mcPDG < 0 ", " mcPDG > 0 "], 'pi+:inRoe': [r'\pi', " mcPDG < 0 ", " mcPDG > 0 "],
82  'h+:inRoe': ['h', " mcPDG < 0 ", " mcPDG > 0 "], 'Lambda0:inRoe': [r'\Lambda', " mcPDG < 0 ", " mcPDG > 0 "]}
83 
84 
85 if not Belle2.FileSystem.findFile('./AsymmetriesInVariablesPlots', True):
86  os.mkdir('./AsymmetriesInVariablesPlots')
87 
88 for (particleList, category, combinerVariable) in ft.eventLevelParticleLists:
89 
90  # if category != "KinLepton":
91  # continue
92 
93  if not Belle2.FileSystem.findFile('./AsymmetriesInVariablesPlots/' + category, True):
94  os.mkdir('./AsymmetriesInVariablesPlots/' + category)
95 
96  if particleList not in identifiersExtraInfosDict and category != 'KaonPion':
97  identifiersExtraInfosDict[particleList] = []
98 
99  methodPrefixEventLevel = "FlavorTagger_" + belleOrBelle2Flag + "_" + weightFiles + 'EventLevel' + category + 'FBDT'
100  treeName = methodPrefixEventLevel + "_tree"
101  targetVariable = 'isRightCategory(' + category + ')'
102 
103  tree = ROOT.TChain(treeName)
104 
105  workingFiles = glob.glob(filesDirectory + '/' + methodPrefixEventLevel + 'sampled*.root')
106  # workingFiles = glob.glob(filesDirectory + '/' + methodPrefixEventLevel + 'sampled1?.root')
107 
108  for iFile in workingFiles:
109  # if Belle2.FileSystem.findFile(workingFile):
110  tree.AddFile(iFile)
111 
112  categoryInputVariables = []
113  trulyUsedInputVariables = []
114  for iVariable in tree.GetListOfBranches():
115 
116  managerVariableName = str(Belle2.invertMakeROOTCompatible(iVariable.GetName()))
117 
118  if managerVariableName in ft.variables[category] or managerVariableName == 'distance' or\
119  managerVariableName == 'z0' or managerVariableName == 'ImpactXY' or\
120  managerVariableName == 'y' or managerVariableName == 'OBoost':
121  if managerVariableName in categoryInputVariables:
122  continue
123 
124  categoryInputVariables.append(managerVariableName)
125  if managerVariableName in ft.variables[category]:
126  allInputVariables.append((category, managerVariableName))
127  trulyUsedInputVariables.append((category, managerVariableName))
128 
129  if managerVariableName not in identifiersExtraInfosDict[particleList] and category != 'KaonPion':
130  identifiersExtraInfosDict[particleList].append(managerVariableName)
131 
132  elif category == 'KaonPion' and managerVariableName not in identifiersExtraInfosKaonPion:
133  identifiersExtraInfosKaonPion.append(managerVariableName)
134 
135  print("The number of variables used in " + category + " is = ", len(trulyUsedInputVariables))
136 
137  if category != 'KaonPion' and category != 'FSC' and category != 'MaximumPstar' and\
138  category != 'FastHadron' and category != 'Lambda':
139  categoryInputVariables.append('BtagToWBosonVariables(recoilMass)')
140 
141  for inputVariable in categoryInputVariables:
142 
143  # if not (inputVariable == 'z0' or inputVariable == 'ImpactXY' or inputVariable == 'y' or inputVariable == 'OBoost'):
144  # continue
145 
146  print(inputVariable)
147 
148  nBins = variablesPlotParamsDict[inputVariable][1]
149  limXmin = variablesPlotParamsDict[inputVariable][2]
150  limXmax = variablesPlotParamsDict[inputVariable][3]
151 
152  if belleOrBelle2 == "Belle2":
153  if inputVariable == 'distance':
154  limXmax = 1.0
155 
156  if inputVariable == 'ImpactXY':
157  limXmax = 0.3
158 
159  if category == "SlowPion":
160  if inputVariable == 'p' or inputVariable == 'useCMSFrame(p)' or\
161  inputVariable == 'pt' or inputVariable == 'useCMSFrame(pt)':
162  nBins = 150
163  limXmax = 1.5
164 
165  if inputVariable == 'distance':
166  nBins = 80
167  limXmax = 2.4
168 
169  if inputVariable == 'ImpactXY':
170  nBins = 80
171  limXmax = 0.8
172 
173  if category == "Lambda":
174  if inputVariable == 'distance':
175  # nBins = 25
176  limXmax = 10
177 
178  if inputVariable == 'chiProb':
179  nBins = 25
180 
181  additionalCond = str()
182 
183  if category == "FastHadron":
184  particleList = 'h+:inRoe'
185 
186  if category == "KinLepton" or category == "IntermediateKinLepton":
187  particleList = 'lepton+:inRoe'
188 
189  negativeHistogram = ROOT.TH1F("negative" + category + str(Belle2.makeROOTCompatible(inputVariable)), "",
190  nBins,
191  limXmin,
192  limXmax)
193  positiveHistogram = ROOT.TH1F("positive" + category + str(Belle2.makeROOTCompatible(inputVariable)), "",
194  nBins,
195  limXmin,
196  limXmax)
197 
198  factorMultiplication = str()
199 
200  if belleOrBelle2 == "Belle2" and ((category != "Lambda" and inputVariable == 'distance') or inputVariable ==
201  'z0' or inputVariable == 'ImpactXY' or inputVariable == 'y' or inputVariable == 'OBoost'):
202  factorMultiplication = "*10 "
203 
204  tree.Draw(variablesPlotParamsDict[inputVariable][0] +
205  factorMultiplication +
206  ">> negative" +
207  category +
208  str(Belle2.makeROOTCompatible(inputVariable)), additionalCond +
209  "abs(" +
210  Belle2.makeROOTCompatible(targetVariable) +
211  ") > 0 && " +
212  particleConditions[particleList][1])
213 
214  tree.Draw(variablesPlotParamsDict[inputVariable][0] +
215  factorMultiplication +
216  ">> positive" +
217  category +
218  str(Belle2.makeROOTCompatible(inputVariable)), additionalCond +
219  "abs(" +
220  Belle2.makeROOTCompatible(targetVariable) +
221  ") > 0 && " +
222  particleConditions[particleList][2])
223 
224  negativeScalingFactor = negativeHistogram.Integral()
225  positiveScalingFactor = positiveHistogram.Integral()
226 
227  if negativeScalingFactor == 0:
228  negativeScalingFactor = 1
229 
230  if positiveScalingFactor == 0:
231  positiveScalingFactor = 1
232 
233  negativeHistogram.Scale(1 / negativeScalingFactor)
234  positiveHistogram.Scale(1 / positiveScalingFactor)
235 
236  negativeArray = np.zeros((negativeHistogram.GetNbinsX(), 2))
237  positiveArray = np.zeros((positiveHistogram.GetNbinsX(), 2))
238 
239  asymmetryArray = np.zeros((positiveHistogram.GetNbinsX(), 2))
240  asymmetryArrayUncertainties = np.zeros((positiveHistogram.GetNbinsX(), 2))
241 
242  for i in range(0, negativeHistogram.GetNbinsX()):
243  negativeArray[i] = np.array([negativeHistogram.GetBinCenter(i + 1), negativeHistogram.GetBinContent(i + 1)])
244  positiveArray[i] = np.array([positiveHistogram.GetBinCenter(i + 1), positiveHistogram.GetBinContent(i + 1)])
245 
246  numerator = float(positiveArray[i][1] - negativeArray[i][1])
247  denominator = float(positiveArray[i][1] + negativeArray[i][1])
248 
249  if denominator == 0:
250  asymmetryArray[i] = np.array([negativeHistogram.GetBinCenter(i + 1), 0])
251  asymmetryArrayUncertainties[i] = np.array([negativeHistogram.GetBinCenter(i + 1), 0])
252  else:
253  asymmetryArray[i] = np.array([negativeHistogram.GetBinCenter(i + 1), float(numerator / denominator)])
254 
255  uncertainty = 2 * math.sqrt((negativeArray[i][1] * negativeHistogram.GetBinError(i + 1))**2 + (
256  positiveArray[i][1] * positiveHistogram.GetBinError(i + 1))**2) / (negativeArray[i][1] + positiveArray[i][1])**2
257 
258  asymmetryArrayUncertainties[i] = np.array([negativeHistogram.GetBinCenter(i + 1), uncertainty])
259 
260  fig1 = plt.figure(1, figsize=(11, 11))
261 
262  # if inputVariable == 'Kid_dEdx' or inputVariable == 'muid_dEdx':
263  # ax1 = plt.axes([0.18, 0.17, 0.75, 0.8])
264  # if inputVariable == 'pi_vs_edEdxid':
265  # ax1 = plt.axes([0.18, 0.187, 0.75, 0.805])
266  # else:
267  ax1 = plt.axes([0.19, 0.37, 0.76, 0.60])
268 
269  # print(negativeArray.shape, negativeHistogram.GetNbinsX(), )
270  # print(negativeArray)
271  ax1.hist(
272  negativeArray[:, 0], weights=negativeArray[:, 1], bins=negativeHistogram.GetNbinsX(),
273  histtype='step', edgecolor='b', linewidth=4.5, linestyle='dashed',
274  label=r'$' + particleConditions[particleList][0] + '^{-} $')
275 
276  ax1.hist(positiveArray[:, 0], weights=positiveArray[:, 1], bins=positiveHistogram.GetNbinsX(),
277  histtype='step', edgecolor='r', linewidth=4, alpha=0.7,
278  label=r'$' + particleConditions[particleList][0] + '^{+} $') # hatch='.',
279 
280  p1, = ax1.plot([], label=r'$' + particleConditions[particleList][0] + '^{-} $',
281  linewidth=5.5, linestyle='dashed', c='b')
282  p2, = ax1.plot([], label=r'$' + particleConditions[particleList][0] +
283  '^{+} $', linewidth=5, linestyle='solid', alpha=0.9, c='r')
284 
285  binWidth = negativeHistogram.GetBinWidth(2)
286 
287  if inputVariable == 'lambdaZError': # or inputVariable == 'ImpactXY' or\
288  # (category != "Lambda" and inputVariable == 'distance'):
289  binWidth = binWidth * 10
290 
291  if inputVariable == 'M':
292  binWidth = binWidth * 1000
293 
294  if category == "Lambda" and inputVariable == 'distance':
295  variablesPlotParamsDict[inputVariable][5] = r"{\rm cm}\, "
296 
297  binWidth = '{:8.2f}'.format(binWidth)
298 
299  legendLocation = 1
300  ax1.set_ylabel(r'${\rm Fraction\hspace{0.25em} of\hspace{0.25em} Events}\, /\, (\, ' + binWidth + r'\, ' +
301  variablesPlotParamsDict[inputVariable][5] + r')$', fontsize=35)
302  if inputVariable == 'extraInfo(isRightCategory(Kaon))' or\
303  inputVariable == 'HighestProbInCat(pi+:inRoe, isRightCategory(SlowPion))':
304  legendLocation = 3
305  ax1.set_yscale('log', nonposy='clip')
306  else:
307  ax1.yaxis.set_major_formatter(FormatStrFormatter(r'$%.2f$'))
308 
309  # ax1.set_ylim(0, 1.4)
310  ax1.set_xlim(limXmin, limXmax)
311 
312  locs, labels = plt.xticks()
313  empty_string_labels = [''] * len(labels)
314  plt.locator_params(axis='x', nbins=len(labels))
315  ax1.set_xticklabels(empty_string_labels)
316  ax1.tick_params(axis='y', labelsize=37)
317 
318  if inputVariable.find('ARICH') != -1 or inputVariable.find('TOP') != -1 or \
319  inputVariable == 'cosTPTO' or inputVariable.find('KLM') != -1 or\
320  inputVariable == 'cosTheta' or inputVariable == 'FSCVariables(cosTPTOFast)' or\
321  inputVariable == 'KaonPionVariables(cosKaonPion)':
322  legendLocation = 2
323 
324  elif inputVariable == 'FSCVariables(SlowFastHaveOpositeCharges)' or \
325  inputVariable == 'KaonPionVariables(HaveOpositeCharges)' or inputVariable == "eid_ECL" or \
326  inputVariable.find('ID') != -1 or inputVariable.find('dEdx') != -1:
327  legendLocation = 9
328 
329  if inputVariable == 'muid_dEdx':
330  if category != 'KinLepton':
331  legendLocation = 8
332 
333  ax1.legend([p1, p2], [r'$' +
334  particleConditions[particleList][0] +
335  '^{-} $', r'$' +
336  particleConditions[particleList][0] +
337  '^{+} $'], prop={'size': 50}, loc=legendLocation, numpoints=1, handlelength=1)
338  ax1.grid(True)
339 
340  ax2 = plt.axes([0.19, 0.15, 0.76, 0.2])
341 
342  # print(asymmetryArrayUncertainties)
343 
344  ax2.errorbar(asymmetryArray[:, 0], asymmetryArray[:, 1], xerr=float(negativeHistogram.GetBinWidth(2) / 2),
345  yerr=asymmetryArrayUncertainties[:, 1], elinewidth=2, mew=2, ecolor='k',
346  fmt='o', mfc='k', mec='k', markersize=6, label=r'${\rm Data}$')
347 
348  ax2.set_ylabel(r'$\frac{f^{+}\; -\;\, f^{-}}{f^{+}\; +\;\, f^{-}}$', fontsize=50, labelpad=20)
349  ax2.yaxis.labelpad = 8
350 
351  plt.yticks([-0.75, -0.5, -0.25, 0, 0.25, 0.5, 0.75],
352  [r'', r'$-0.5$', r'', r'$0.0$', r'', r'$0.5$', r''], rotation=0, size=25)
353  ax2.set_ylim(-0.75, 0.75)
354  ax2.xaxis.grid(True) # linestyle='--'
355  plt.axhline(y=0.5, linewidth=2, color='tab:gray', linestyle='-')
356  plt.axhline(y=0.25, linewidth=2, color='tab:gray', linestyle='-')
357  plt.axhline(y=0, linewidth=2, color='tab:gray', linestyle='-')
358  plt.axhline(y=-0.25, linewidth=2, color='tab:gray', linestyle='-')
359  plt.axhline(y=-0.5, linewidth=2, color='tab:gray', linestyle='-')
360 
361  xLabel = variablesPlotParamsDict[inputVariable][4]
362  if category == "FSC":
363  if inputVariable == 'cosTPTO':
364  xLabel = r'$\vert\cos{\theta^*_{\rm T, Slow}}\vert$'
365  if inputVariable == 'useCMSFrame(p)':
366  xLabel = r'$p^*_{\rm Slow}\ [{\rm GeV}/c]$'
367  if category == 'Lambda':
368  if inputVariable == 'useCMSFrame(p)':
369  xLabel = r'$p^*_{\Lambda}\ [{\rm GeV}/c]$'
370  if inputVariable == 'p':
371  xLabel = r'$p_{\Lambda}\ [{\rm GeV}/c]$'
372 
373  if inputVariable == 'pi_vs_edEdxid':
374  ax2.xaxis.labelpad = 5
375  else:
376  ax2.xaxis.labelpad = 15
377  plt.locator_params(axis='x', nbins=len(labels))
378 
379  ax2.set_xlim(limXmin, limXmax)
380  ax2.tick_params(axis='x', labelsize=40)
381  ax2.set_xlabel(xLabel, fontsize=50)
382  plt.savefig('./AsymmetriesInVariablesPlots/' + category + '/' + category +
383  "_" + str(Belle2.makeROOTCompatible(inputVariable)) + '.pdf')
384  fig1.clear()
385 
386  negativeHistogram.Delete()
387  positiveHistogram.Delete()
388 
389 totalNumberOfVariables = 0
390 
391 for category in ft.variables:
392  totalNumberOfVariables += len(ft.variables[category])
393 
394 print("Total number of variables = ", totalNumberOfVariables)
395 
396 totalNumberOfCalculatedVariables = len(identifiersExtraInfosKaonPion)
397 
398 print("Calculations for Kaon-Pion Category = ", totalNumberOfCalculatedVariables)
399 
400 print("Variables per particle list:")
401 for particleList in identifiersExtraInfosDict:
402  print(particleList)
403  print(identifiersExtraInfosDict[particleList])
404  totalNumberOfCalculatedVariables += len(identifiersExtraInfosDict[particleList])
405 
406 print("Total number of calculated variables = ", totalNumberOfCalculatedVariables)
B2A714-DeepContinuumSuppression_MVAModel.batch_generator.len
len
Number of events.
Definition: B2A714-DeepContinuumSuppression_MVAModel.py:70
Belle2::invertMakeROOTCompatible
std::string invertMakeROOTCompatible(std::string str)
Invert makeROOTCompatible operation.
Definition: MakeROOTCompatible.cc:89
Belle2::makeROOTCompatible
std::string makeROOTCompatible(std::string str)
Remove special characters that ROOT dislikes in branch names, e.g.
Definition: MakeROOTCompatible.cc:74
Belle2::FileSystem::findFile
static std::string findFile(const std::string &path, bool silent=false)
Search for given file or directory in local or central release directory, and return absolute path if...
Definition: FileSystem.cc:147