Belle II Software  release-05-01-25
kaons_pdf.py
1 # !/usr/bin/env python3
2 
3 """ kaons_pdf.py : construct the pdf for kaon ECL charged PID. """
4 
5 __author__ = "Marco Milesi"
6 __email__ = "marco.milesi@unimelb.edu.au"
7 __maintainer__ = "Marco Milesi"
8 
9 import os
10 import ROOT
11 from plotting_utils import draw_plots, get_ndf
12 
13 
14 def fit_kaon_eop(**kwargs):
15 
16  append = "anti" if kwargs["charge"] < 0 else ""
17 
18  idx_p = kwargs["idx_p"]
19  idx_th = kwargs["idx_theta"]
20 
21  # Set the PDF paarameter ranges ((start), min, max).
22  if kwargs["pmin"] < 5.5e3:
23  # Bifurgaus for the peak, gaus0 for the right kink, gaus1 for the long right tail.
24  eop_range = (0.0, 0.8)
25  bifurgaus_mu_range = (0.03, 0.01, 0.1)
26  bifurgaus_sigmaL_range = (0.01, 0.001, 0.1)
27  bifurgaus_sigmaR_range = (0.01, 0.001, 0.1)
28  bifurgaus_frac_range = (0.8, 1.0)
29  gaus0_mu_range = (0.07, 0.05, 0.2)
30  gaus0_sigma_range = (0.015, 0.001, 0.05)
31  gaus1_mu_range = (0.4, 0.1, 0.8)
32  gaus1_sigma_range = (0.1, 0.01, 0.7)
33  gaus1_frac_range = (0.0, 1.0)
34  if kwargs["pmin"] < 4e3:
35  eop_range = (0.0, 0.8)
36  bifurgaus_mu_range = (0.05, 0.01, 0.1)
37  bifurgaus_sigmaL_range = (0.01, 0.001, 0.05)
38  bifurgaus_sigmaR_range = (0.02, 0.001, 0.1)
39  bifurgaus_frac_range = (0.6, 1.0)
40  gaus0_mu_range = (0.075, 0.05, 0.2)
41  gaus0_sigma_range = (0.05, 0.01, 0.1)
42  gaus1_mu_range = (0.4, 0.1, 0.8)
43  gaus1_sigma_range = (0.1, 0.01, 0.7)
44  gaus1_frac_range = (0.0, 1.0)
45  if kwargs["pmin"] < 3e3:
46  eop_range = (0.0, 1.0)
47  bifurgaus_mu_range = (0.065, 0.01, 0.15)
48  bifurgaus_sigmaL_range = (0.01, 0.001, 0.05)
49  bifurgaus_sigmaR_range = (0.02, 0.001, 0.1)
50  bifurgaus_frac_range = (0.6, 1.0)
51  gaus0_mu_range = (0.1, 0.05, 0.2)
52  gaus0_sigma_range = (0.05, 0.01, 0.1)
53  gaus1_mu_range = (0.4, 0.1, 0.8)
54  gaus1_sigma_range = (0.1, 0.01, 0.7)
55  gaus1_frac_range = (0.0, 1.0)
56  if kwargs["pmin"] < 2e3:
57  eop_range = (0.0, 1.2)
58  bifurgaus_mu_range = (0.12, 0.01, 0.2)
59  bifurgaus_sigmaL_range = (0.01, 0.001, 0.05)
60  bifurgaus_sigmaR_range = (0.02, 0.001, 0.1)
61  bifurgaus_frac_range = (0.6, 1.0)
62  gaus0_mu_range = (0.15, 0.05, 0.4)
63  gaus0_sigma_range = (0.1, 0.01, 0.2)
64  gaus1_mu_range = (0.4, 0.1, 0.8)
65  gaus1_sigma_range = (0.1, 0.01, 0.7)
66  gaus1_frac_range = (0.0, 1.0)
67  if kwargs["pmin"] < 1.5e3:
68  eop_range = (0.0, 1.2)
69  bifurgaus_mu_range = (0.13, 0.1, 0.3)
70  bifurgaus_sigmaL_range = (0.01, 0.001, 0.1)
71  bifurgaus_sigmaR_range = (0.05, 0.001, 0.1)
72  bifurgaus_frac_range = (0.5, 1.0)
73  gaus0_mu_range = (0.25, 0.2, 0.5)
74  gaus0_sigma_range = (0.1, 0.01, 0.2)
75  gaus1_mu_range = (0.4, 0.1, 0.8)
76  gaus1_sigma_range = (0.1, 0.01, 0.7)
77  gaus1_frac_range = (0.0, 1.0)
78  if kwargs["pmin"] < 1e3:
79  eop_range = (0.0, 1.2)
80  bifurgaus_mu_range = (0.2, 0.1, 0.4)
81  bifurgaus_sigmaL_range = (0.01, 0.001, 0.1)
82  bifurgaus_sigmaR_range = (0.05, 0.001, 0.1)
83  bifurgaus_frac_range = (0.5, 1.0)
84  gaus0_mu_range = (0.5, 0.35, 0.65)
85  gaus0_sigma_range = (0.1, 0.01, 0.2)
86  gaus1_mu_range = (0.5, 0.35, 0.8)
87  gaus1_sigma_range = (0.1, 0.01, 0.7)
88  gaus1_frac_range = (0.0, 1.0)
89  if kwargs["pmin"] < 0.75e3:
90  # No more mid-E/p kink; now gaus0 describes the feature at low E/p
91  eop_range = (0.0, 1.4)
92  bifurgaus_mu_range = (0.35, 0.15, 0.5) if kwargs["charge"] > 0 else (0.5, 0.3, 0.7)
93  bifurgaus_sigmaL_range = (0.1, 0.001, 0.2)
94  bifurgaus_sigmaR_range = (0.1, 0.001, 0.2)
95  bifurgaus_frac_range = (0.0, 1.0)
96  gaus0_mu_range = (0.05, 0.0, 0.1)
97  gaus0_sigma_range = (0.025, 0.0, 0.05)
98  gaus1_mu_range = (0.8, 0.1, 1.0) if kwargs["charge"] > 0 else (1.0, 0.7, 1.4)
99  gaus1_sigma_range = (0.2, 0.01, 0.7)
100  gaus1_frac_range = (0.0, 1.0)
101  if kwargs["pmin"] < 0.5e3:
102  eop_range = (0.0, 1.5)
103  bifurgaus_mu_range = (0.35, 0.15, 0.5) if kwargs["charge"] > 0 else (0.5, 0.3, 0.7)
104  bifurgaus_sigmaL_range = (0.1, 0.001, 0.2)
105  bifurgaus_sigmaR_range = (0.1, 0.001, 0.2)
106  bifurgaus_frac_range = (0.0, 1.0)
107  gaus0_mu_range = (0.05, 0.0, 0.1)
108  gaus0_sigma_range = (0.025, 0.0, 0.05)
109  gaus1_mu_range = (0.8, 0.1, 1.0) if kwargs["charge"] > 0 else (1.0, 0.7, 1.4)
110  gaus1_sigma_range = (0.2, 0.01, 0.7)
111  gaus1_frac_range = (0.0, 1.0)
112  if kwargs["pmin"] < 0.4e3:
113  eop_range = (0.0, 1.5)
114  bifurgaus_mu_range = (0.45, 0.15, 0.6) if kwargs["charge"] > 0 else (0.35, 0.2, 0.7)
115  bifurgaus_sigmaL_range = (0.1, 0.001, 0.2)
116  bifurgaus_sigmaR_range = (0.1, 0.001, 0.2)
117  bifurgaus_frac_range = (0.0, 1.0)
118  gaus0_mu_range = (0.05, 0.0, 0.2)
119  gaus0_sigma_range = (0.025, 0.0, 0.05)
120  gaus1_mu_range = (0.8, 0.1, 1.0)
121  gaus1_sigma_range = (0.2, 0.01, 0.7)
122  gaus1_frac_range = (0.0, 1.0)
123 
124  # Open file w/ input E/p distribution.
125  inpath = "{0}/pdg{1}321.root".format(kwargs["inputpath"], append)
126  infile = ROOT.TFile(inpath)
127 
128  # Create variable to fit and its distribution in data.
129  eopvar = ROOT.RooRealVar("eop", "E/p", eop_range[0], eop_range[1])
130  eophist_data = infile.Get("h_Eop_{0}_{1}".format(idx_p - 1, idx_th - 1))
131  eopdata = ROOT.RooDataHist("eopdata", "eopdata", ROOT.RooArgList(eopvar), ROOT.RooFit.Import(eophist_data))
132 
133  # Create plot for the fit.
134  pm = "+" if kwargs["charge"] > 0 else "-"
135  particle = "K^{{{0}}}".format(pm)
136  p_range = "{0:.2f} < p_{{lab}} #leq {1:.2f} [GeV/c]".format(kwargs["pmin"] / 1e3, kwargs["pmax"] / 1e3)
137  theta_range = "{0:.1f} < #theta_{{lab}} #leq {1:.1f} [deg]".format(kwargs["thetamin"], kwargs["thetamax"])
138  title = "{0}, {1}, {2}".format(particle, p_range, theta_range)
139  frame1 = eopvar.frame(ROOT.RooFit.Title(title))
140  eopdata.plotOn(frame1, ROOT.RooFit.DataError(ROOT.RooAbsData.SumW2))
141 
142  # Create PDF to fit.
143  bifurgaus_mu = ROOT.RooRealVar("#mu_{bg}",
144  "mean of bifurcated gaussian",
145  bifurgaus_mu_range[0],
146  bifurgaus_mu_range[1],
147  bifurgaus_mu_range[2])
148  bifurgaus_sigmal = ROOT.RooRealVar("#sigma_{bg}_{L}",
149  "widthL of bifurcated gaussian",
150  bifurgaus_sigmaL_range[0],
151  bifurgaus_sigmaL_range[1],
152  bifurgaus_sigmaL_range[2])
153  bifurgaus_sigmar = ROOT.RooRealVar("#sigma_{bg}_{R}",
154  "widthR of bifurcated gaussian",
155  bifurgaus_sigmaR_range[0],
156  bifurgaus_sigmaR_range[1],
157  bifurgaus_sigmaR_range[2])
158  bifurgaus_frac = ROOT.RooRealVar("frac_{bg}",
159  "bifurcated gaussian fraction",
160  bifurgaus_frac_range[0],
161  bifurgaus_frac_range[1])
162  bifurgaus = ROOT.RooBifurGauss("bifurgaus", "bifurcated gaussian", eopvar, bifurgaus_mu, bifurgaus_sigmal, bifurgaus_sigmar)
163 
164  gaus0_mean = ROOT.RooRealVar("#mu_{g0}",
165  "mean of gaussian",
166  gaus0_mu_range[0],
167  gaus0_mu_range[1],
168  gaus0_mu_range[2])
169  gaus0_sigma = ROOT.RooRealVar("#sigma_{g0}",
170  "width of gaussian",
171  gaus0_sigma_range[0],
172  gaus0_sigma_range[1],
173  gaus0_sigma_range[2])
174  gaus0 = ROOT.RooGaussian("gaus0", "gaussian", eopvar, gaus0_mean, gaus0_sigma)
175 
176  pdf0 = ROOT.RooAddPdf("pdf0", "bifurcated gaussian + gaussian", bifurgaus, gaus0, bifurgaus_frac)
177 
178  gaus1_mean = ROOT.RooRealVar("#mu_{g1}",
179  "mean of gaussian",
180  gaus1_mu_range[0],
181  gaus1_mu_range[1],
182  gaus1_mu_range[2])
183  gaus1_sigma = ROOT.RooRealVar("#sigma_{g1}",
184  "width of gaussian",
185  gaus1_sigma_range[0],
186  gaus1_sigma_range[1],
187  gaus1_sigma_range[2])
188  gaus1_frac = ROOT.RooRealVar("frac_{g1}",
189  "gaussian fraction",
190  gaus1_frac_range[0],
191  gaus1_frac_range[1])
192  gaus1 = ROOT.RooGaussian("gaus1", "gaussian", eopvar, gaus1_mean, gaus1_sigma)
193 
194  pdf = ROOT.RooAddPdf("pdf", "bifurcated gaussian + gaussian + gaussian", pdf0, gaus1, gaus1_frac)
195 
196  # Ok, fit!
197  fitres = pdf.fitTo(eopdata, ROOT.RooFit.Save(), ROOT.RooFit.PrintEvalErrors(-1))
198 
199  # Extract the normalised post-fit PDF as TF1.
200  ral1 = ROOT.RooArgList(eopvar)
201  ral2 = ROOT.RooArgList(pdf.getParameters(ROOT.RooArgSet(eopvar)))
202  ras = ROOT.RooArgSet(eopvar)
203  pdffunc = pdf.asTF(ral1, ral2, ras)
204 
205  # Add the PDFs and the fitted parameters to the plot.
206  pdf.plotOn(frame1,
207  ROOT.RooFit.LineColor(ROOT.kBlack),
208  ROOT.RooFit.LineWidth(2),
209  ROOT.RooFit.MoveToBack())
210  pdf.plotOn(frame1,
211  ROOT.RooFit.Components("bifurgaus"),
212  ROOT.RooFit.LineStyle(ROOT.kDashed),
213  ROOT.RooFit.LineColor(ROOT.kBlue),
214  ROOT.RooFit.LineWidth(2),
215  ROOT.RooFit.MoveToBack())
216  pdf.plotOn(frame1,
217  ROOT.RooFit.Components("gaus0"),
218  ROOT.RooFit.LineStyle(ROOT.kSolid),
219  ROOT.RooFit.LineColor(ROOT.kAzure + 10),
220  ROOT.RooFit.LineWidth(2),
221  ROOT.RooFit.MoveToBack())
222  pdf.plotOn(frame1,
223  ROOT.RooFit.Components("gaus1"),
224  ROOT.RooFit.LineStyle(ROOT.kDotted),
225  ROOT.RooFit.LineColor(ROOT.kCyan - 3),
226  ROOT.RooFit.LineWidth(2),
227  ROOT.RooFit.MoveToBack())
228  pdf.paramOn(frame1, ROOT.RooFit.Layout(0.55, 0.9, 0.75))
229 
230  # Create plot for fit residuals.
231  frame2 = eopvar.frame(ROOT.RooFit.Title("Residuals"))
232  frame2.addPlotable(frame1.residHist(), "P")
233 
234  # Check the goodness of fit.
235  ndf = get_ndf(eopdata, fitres)
236  print("X^2/ndf = {0:.3f} (ndf = {1})".format(frame1.chiSquare(), ndf))
237 
238  # Draw the plots and save them.
239  if kwargs["outputplots"]:
240  plotargs = dict(kwargs, frame1=frame1.Clone(), frame2=frame2.Clone(), ndf=ndf, ndiv=520, append=append)
241  draw_plots(**plotargs)
242 
243  pdfname = "pdf_EoP_{0}_{1}".format(idx_p, idx_th)
244  return pdffunc.Clone(pdfname) # Why cloning? A: https://root-forum.cern.ch/t/returning-th1f-from-function/17213/2