20 from ROOT
import Belle2
22 import flavorTagger
as ft
23 from defaultEvaluationParameters
import r_subsample, r_size, categories
24 from inputVariablesPlots
import variablesPlotParamsDict
26 from array
import array
29 import matplotlib
as mpl
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
41 if len(sys.argv) != 5:
42 sys.exit(
"Must provide 4 arguments: [Belle or Belle2] [BGx0 or BGx1] [training] [workingDirectory]"
45 belleOrBelle2 = sys.argv[1]
46 MCtype = str(sys.argv[2])
47 decayChannelTrainedOn = str(sys.argv[3])
48 workingDirectory = sys.argv[4]
50 filesDirectory = workingDirectory +
'/FlavorTagging/TrainedMethods'
52 if decayChannelTrainedOn ==
'JPsiKs':
53 decayChannelTrainedOn =
'JpsiKs_mu'
55 weightFiles =
'B2' + decayChannelTrainedOn + MCtype
57 ROOT.TH1.SetDefaultSumw2()
59 allInputVariables = []
61 belleOrBelle2Flag = belleOrBelle2
63 identifiersExtraInfosDict = dict()
64 identifiersExtraInfosKaonPion = []
73 if belleOrBelle2 ==
"Belle":
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 "]}
86 os.mkdir(
'./AsymmetriesInVariablesPlots')
88 for (particleList, category, combinerVariable)
in ft.eventLevelParticleLists:
94 os.mkdir(
'./AsymmetriesInVariablesPlots/' + category)
96 if particleList
not in identifiersExtraInfosDict
and category !=
'KaonPion':
97 identifiersExtraInfosDict[particleList] = []
99 methodPrefixEventLevel =
"FlavorTagger_" + belleOrBelle2Flag +
"_" + weightFiles +
'EventLevel' + category +
'FBDT'
100 treeName = methodPrefixEventLevel +
"_tree"
101 targetVariable =
'isRightCategory(' + category +
')'
103 tree = ROOT.TChain(treeName)
105 workingFiles = glob.glob(filesDirectory +
'/' + methodPrefixEventLevel +
'sampled*.root')
108 for iFile
in workingFiles:
112 categoryInputVariables = []
113 trulyUsedInputVariables = []
114 for iVariable
in tree.GetListOfBranches():
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:
124 categoryInputVariables.append(managerVariableName)
125 if managerVariableName
in ft.variables[category]:
126 allInputVariables.append((category, managerVariableName))
127 trulyUsedInputVariables.append((category, managerVariableName))
129 if managerVariableName
not in identifiersExtraInfosDict[particleList]
and category !=
'KaonPion':
130 identifiersExtraInfosDict[particleList].append(managerVariableName)
132 elif category ==
'KaonPion' and managerVariableName
not in identifiersExtraInfosKaonPion:
133 identifiersExtraInfosKaonPion.append(managerVariableName)
135 print(
"The number of variables used in " + category +
" is = ",
len(trulyUsedInputVariables))
137 if category !=
'KaonPion' and category !=
'FSC' and category !=
'MaximumPstar' and\
138 category !=
'FastHadron' and category !=
'Lambda':
139 categoryInputVariables.append(
'BtagToWBosonVariables(recoilMass)')
141 for inputVariable
in categoryInputVariables:
148 nBins = variablesPlotParamsDict[inputVariable][1]
149 limXmin = variablesPlotParamsDict[inputVariable][2]
150 limXmax = variablesPlotParamsDict[inputVariable][3]
152 if belleOrBelle2 ==
"Belle2":
153 if inputVariable ==
'distance':
156 if inputVariable ==
'ImpactXY':
159 if category ==
"SlowPion":
160 if inputVariable ==
'p' or inputVariable ==
'useCMSFrame(p)' or\
161 inputVariable ==
'pt' or inputVariable ==
'useCMSFrame(pt)':
165 if inputVariable ==
'distance':
169 if inputVariable ==
'ImpactXY':
173 if category ==
"Lambda":
174 if inputVariable ==
'distance':
178 if inputVariable ==
'chiProb':
181 additionalCond = str()
183 if category ==
"FastHadron":
184 particleList =
'h+:inRoe'
186 if category ==
"KinLepton" or category ==
"IntermediateKinLepton":
187 particleList =
'lepton+:inRoe'
198 factorMultiplication = str()
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 "
204 tree.Draw(variablesPlotParamsDict[inputVariable][0] +
205 factorMultiplication +
212 particleConditions[particleList][1])
214 tree.Draw(variablesPlotParamsDict[inputVariable][0] +
215 factorMultiplication +
222 particleConditions[particleList][2])
224 negativeScalingFactor = negativeHistogram.Integral()
225 positiveScalingFactor = positiveHistogram.Integral()
227 if negativeScalingFactor == 0:
228 negativeScalingFactor = 1
230 if positiveScalingFactor == 0:
231 positiveScalingFactor = 1
233 negativeHistogram.Scale(1 / negativeScalingFactor)
234 positiveHistogram.Scale(1 / positiveScalingFactor)
236 negativeArray = np.zeros((negativeHistogram.GetNbinsX(), 2))
237 positiveArray = np.zeros((positiveHistogram.GetNbinsX(), 2))
239 asymmetryArray = np.zeros((positiveHistogram.GetNbinsX(), 2))
240 asymmetryArrayUncertainties = np.zeros((positiveHistogram.GetNbinsX(), 2))
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)])
246 numerator = float(positiveArray[i][1] - negativeArray[i][1])
247 denominator = float(positiveArray[i][1] + negativeArray[i][1])
250 asymmetryArray[i] = np.array([negativeHistogram.GetBinCenter(i + 1), 0])
251 asymmetryArrayUncertainties[i] = np.array([negativeHistogram.GetBinCenter(i + 1), 0])
253 asymmetryArray[i] = np.array([negativeHistogram.GetBinCenter(i + 1), float(numerator / denominator)])
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
258 asymmetryArrayUncertainties[i] = np.array([negativeHistogram.GetBinCenter(i + 1), uncertainty])
260 fig1 = plt.figure(1, figsize=(11, 11))
267 ax1 = plt.axes([0.19, 0.37, 0.76, 0.60])
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] +
'^{-} $')
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] +
'^{+} $')
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')
285 binWidth = negativeHistogram.GetBinWidth(2)
287 if inputVariable ==
'lambdaZError':
289 binWidth = binWidth * 10
291 if inputVariable ==
'M':
292 binWidth = binWidth * 1000
294 if category ==
"Lambda" and inputVariable ==
'distance':
295 variablesPlotParamsDict[inputVariable][5] =
r"{\rm cm}\, "
297 binWidth =
'{:8.2f}'.format(binWidth)
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))':
305 ax1.set_yscale(
'log', nonposy=
'clip')
307 ax1.yaxis.set_major_formatter(FormatStrFormatter(
r'$%.2f$'))
310 ax1.set_xlim(limXmin, limXmax)
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)
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)':
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:
329 if inputVariable ==
'muid_dEdx':
330 if category !=
'KinLepton':
333 ax1.legend([p1, p2], [
r'$' +
334 particleConditions[particleList][0] +
336 particleConditions[particleList][0] +
337 '^{+} $'], prop={
'size': 50}, loc=legendLocation, numpoints=1, handlelength=1)
340 ax2 = plt.axes([0.19, 0.15, 0.76, 0.2])
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}$')
348 ax2.set_ylabel(
r'$\frac{f^{+}\; -\;\, f^{-}}{f^{+}\; +\;\, f^{-}}$', fontsize=50, labelpad=20)
349 ax2.yaxis.labelpad = 8
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)
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=
'-')
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]$'
373 if inputVariable ==
'pi_vs_edEdxid':
374 ax2.xaxis.labelpad = 5
376 ax2.xaxis.labelpad = 15
377 plt.locator_params(axis=
'x', nbins=
len(labels))
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 +
386 negativeHistogram.Delete()
387 positiveHistogram.Delete()
389 totalNumberOfVariables = 0
391 for category
in ft.variables:
392 totalNumberOfVariables +=
len(ft.variables[category])
394 print(
"Total number of variables = ", totalNumberOfVariables)
396 totalNumberOfCalculatedVariables =
len(identifiersExtraInfosKaonPion)
398 print(
"Calculations for Kaon-Pion Category = ", totalNumberOfCalculatedVariables)
400 print(
"Variables per particle list:")
401 for particleList
in identifiersExtraInfosDict:
403 print(identifiersExtraInfosDict[particleList])
404 totalNumberOfCalculatedVariables +=
len(identifiersExtraInfosDict[particleList])
406 print(
"Total number of calculated variables = ", totalNumberOfCalculatedVariables)