Belle II Software  release-05-01-25
pions_pdf.py
1 # !/usr/bin/env python3
2 
3 """ pions_pdf.py : construct the pdf for pion 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_pion_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.001, 0.0, 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.03, 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.1, 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.07, 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.1, 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.2, 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"] < 1e3:
68  eop_range = (0.0, 1.2)
69  bifurgaus_mu_range = (0.2, 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.3, 1.0)
73  gaus0_mu_range = (0.35, 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"] < 0.75e3:
79  # No more mid-E/p kink; now gaus0 describes the feature at low E/p
80  eop_range = (0.0, 1.2)
81  bifurgaus_mu_range = (0.3, 0.15, 0.45)
82  bifurgaus_sigmaL_range = (0.01, 0.001, 0.1)
83  bifurgaus_sigmaR_range = (0.05, 0.001, 0.1)
84  bifurgaus_frac_range = (0.0, 1.0)
85  gaus0_mu_range = (0.05, 0.0, 0.1)
86  gaus0_sigma_range = (0.025, 0.0, 0.05)
87  gaus1_mu_range = (0.6, 0.1, 0.8)
88  gaus1_sigma_range = (0.2, 0.01, 0.7)
89  gaus1_frac_range = (0.0, 1.0)
90  if kwargs["pmin"] < 0.5e3:
91  eop_range = (0.0, 1.2)
92  bifurgaus_mu_range = (0.4, 0.2, 0.6)
93  bifurgaus_sigmaL_range = (0.1, 0.001, 0.3)
94  bifurgaus_sigmaR_range = (0.2, 0.001, 0.4)
95  bifurgaus_frac_range = (0.0, 1.0)
96  gaus0_mu_range = (0.05, 0.0, 0.2)
97  gaus0_sigma_range = (0.025, 0.0, 0.1)
98  gaus1_mu_range = (0.6, 0.1, 0.8)
99  gaus1_sigma_range = (0.2, 0.01, 0.7)
100  gaus1_frac_range = (0.0, 1.0)
101  if kwargs["pmin"] < 0.4e3:
102  eop_range = (0.0, 1.2)
103  bifurgaus_mu_range = (0.4, 0.2, 0.8)
104  bifurgaus_sigmaL_range = (0.2, 0.001, 0.4)
105  bifurgaus_sigmaR_range = (0.2, 0.001, 0.4)
106  bifurgaus_frac_range = (0.0, 1.0)
107  gaus0_mu_range = (0.1, 0.0, 0.2)
108  gaus0_sigma_range = (0.05, 0.0, 0.1)
109  gaus1_mu_range = (0.6, 0.1, 0.8)
110  gaus1_sigma_range = (0.2, 0.01, 0.7)
111  gaus1_frac_range = (0.0, 1.0)
112 
113  # Open file w/ input E/p distribution.
114  inpath = "{0}/pdg{1}211.root".format(kwargs["inputpath"], append)
115  infile = ROOT.TFile(inpath)
116 
117  # Create variable to fit and its distribution in data.
118  eopvar = ROOT.RooRealVar("eop", "E/p", eop_range[0], eop_range[1])
119  eophist_data = infile.Get("h_Eop_{0}_{1}".format(idx_p - 1, idx_th - 1))
120  eopdata = ROOT.RooDataHist("eopdata", "eopdata", ROOT.RooArgList(eopvar), ROOT.RooFit.Import(eophist_data))
121 
122  # Create plot for the fit.
123  pm = "+" if kwargs["charge"] > 0 else "-"
124  particle = "#pi^{{{0}}}".format(pm)
125  p_range = "{0:.2f} < p_{{lab}} #leq {1:.2f} [GeV/c]".format(kwargs["pmin"] / 1e3, kwargs["pmax"] / 1e3)
126  theta_range = "{0:.1f} < #theta_{{lab}} #leq {1:.1f} [deg]".format(kwargs["thetamin"], kwargs["thetamax"])
127  title = "{0}, {1}, {2}".format(particle, p_range, theta_range)
128  frame1 = eopvar.frame(ROOT.RooFit.Title(title))
129  eopdata.plotOn(frame1, ROOT.RooFit.DataError(ROOT.RooAbsData.SumW2))
130 
131  # Create PDF to fit.
132  bifurgaus_mu = ROOT.RooRealVar("#mu_{bg}",
133  "mean of bifurcated gaussian",
134  bifurgaus_mu_range[0],
135  bifurgaus_mu_range[1],
136  bifurgaus_mu_range[2])
137  bifurgaus_sigmal = ROOT.RooRealVar("#sigma_{bg}_{L}",
138  "widthL of bifurcated gaussian",
139  bifurgaus_sigmaL_range[0],
140  bifurgaus_sigmaL_range[1],
141  bifurgaus_sigmaL_range[2])
142  bifurgaus_sigmar = ROOT.RooRealVar("#sigma_{bg}_{R}",
143  "widthR of bifurcated gaussian",
144  bifurgaus_sigmaR_range[0],
145  bifurgaus_sigmaR_range[1],
146  bifurgaus_sigmaR_range[2])
147  bifurgaus_frac = ROOT.RooRealVar("frac_{bg}",
148  "bifurcated gaussian fraction",
149  bifurgaus_frac_range[0],
150  bifurgaus_frac_range[1])
151  bifurgaus = ROOT.RooBifurGauss("bifurgaus", "bifurcated gaussian", eopvar, bifurgaus_mu, bifurgaus_sigmal, bifurgaus_sigmar)
152 
153  gaus0_mean = ROOT.RooRealVar("#mu_{g0}",
154  "mean of gaussian",
155  gaus0_mu_range[0],
156  gaus0_mu_range[1],
157  gaus0_mu_range[2])
158  gaus0_sigma = ROOT.RooRealVar("#sigma_{g0}",
159  "width of gaussian",
160  gaus0_sigma_range[0],
161  gaus0_sigma_range[1],
162  gaus0_sigma_range[2])
163  gaus0 = ROOT.RooGaussian("gaus0", "gaussian", eopvar, gaus0_mean, gaus0_sigma)
164 
165  pdf0 = ROOT.RooAddPdf("pdf0", "bifurcated gaussian + gaussian", bifurgaus, gaus0, bifurgaus_frac)
166 
167  gaus1_mean = ROOT.RooRealVar("#mu_{g1}",
168  "mean of gaussian",
169  gaus1_mu_range[0],
170  gaus1_mu_range[1],
171  gaus1_mu_range[2])
172  gaus1_sigma = ROOT.RooRealVar("#sigma_{g1}",
173  "width of gaussian",
174  gaus1_sigma_range[0],
175  gaus1_sigma_range[1],
176  gaus1_sigma_range[2])
177  gaus1_frac = ROOT.RooRealVar("frac_{g1}",
178  "gaussian fraction",
179  gaus1_frac_range[0],
180  gaus1_frac_range[1])
181  gaus1 = ROOT.RooGaussian("gaus1", "gaussian", eopvar, gaus1_mean, gaus1_sigma)
182 
183  pdf = ROOT.RooAddPdf("pdf", "bifurcated gaussian + gaussian + gaussian", pdf0, gaus1, gaus1_frac)
184 
185  # Ok, fit!
186  fitres = pdf.fitTo(eopdata, ROOT.RooFit.Save(), ROOT.RooFit.PrintEvalErrors(-1))
187 
188  # Extract the normalised post-fit PDF as TF1.
189  ral1 = ROOT.RooArgList(eopvar)
190  ral2 = ROOT.RooArgList(pdf.getParameters(ROOT.RooArgSet(eopvar)))
191  ras = ROOT.RooArgSet(eopvar)
192  pdffunc = pdf.asTF(ral1, ral2, ras)
193 
194  # Add the PDFs and the fitted parameters to the plot.
195  pdf.plotOn(frame1,
196  ROOT.RooFit.LineColor(ROOT.kBlack),
197  ROOT.RooFit.LineWidth(2),
198  ROOT.RooFit.MoveToBack())
199  pdf.plotOn(frame1,
200  ROOT.RooFit.Components("bifurgaus"),
201  ROOT.RooFit.LineStyle(ROOT.kDashed),
202  ROOT.RooFit.LineColor(ROOT.kGray),
203  ROOT.RooFit.LineWidth(2),
204  ROOT.RooFit.MoveToBack())
205  pdf.plotOn(frame1,
206  ROOT.RooFit.Components("gaus0"),
207  ROOT.RooFit.LineStyle(ROOT.kSolid),
208  ROOT.RooFit.LineColor(ROOT.kGray + 1),
209  ROOT.RooFit.LineWidth(2),
210  ROOT.RooFit.MoveToBack())
211  pdf.plotOn(frame1,
212  ROOT.RooFit.Components("gaus1"),
213  ROOT.RooFit.LineStyle(ROOT.kDotted),
214  ROOT.RooFit.LineColor(ROOT.kGray + 2),
215  ROOT.RooFit.LineWidth(2),
216  ROOT.RooFit.MoveToBack())
217  pdf.paramOn(frame1, ROOT.RooFit.Layout(0.55, 0.9, 0.75))
218 
219  # Create plot for fit residuals.
220  frame2 = eopvar.frame(ROOT.RooFit.Title("Residuals"))
221  frame2.addPlotable(frame1.residHist(), "P")
222 
223  # Check the goodness of fit.
224  ndf = get_ndf(eopdata, fitres)
225  print("X^2/ndf = {0:.3f} (ndf = {1})".format(frame1.chiSquare(), ndf))
226 
227  # Draw the plots and save them.
228  if kwargs["outputplots"]:
229  plotargs = dict(kwargs, frame1=frame1.Clone(), frame2=frame2.Clone(), ndf=ndf, ndiv=520, append=append)
230  draw_plots(**plotargs)
231 
232  pdfname = "pdf_EoP_{0}_{1}".format(idx_p, idx_th)
233  return pdffunc.Clone(pdfname) # Why cloning? A: https://root-forum.cern.ch/t/returning-th1f-from-function/17213/2