StrangenessTrackingQC/task/strangeness_tracking_eff.py at main · fmazzasc/StrangenessTrackingQC · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
import ROOT
import numpy as np

def computeEfficiency(gen_hist, rec_hist, name, rebin=0):
    if rebin > 1:
        gen_hist.Rebin(rebin)
        rec_hist.Rebin(rebin)
    eff_hist = gen_hist.Clone(name)
    eff_hist.Reset()
    eff_hist.GetYaxis().SetTitle(r'Tracked / Found')
    eff_hist.GetYaxis().SetRangeUser(0., 1.1)
    for iPt in range(1, rec_hist.GetNbinsX() + 1):
        gen_val = gen_hist.GetBinContent(iPt)
        if gen_val < 1e-24:
            continue
        rec_val = rec_hist.GetBinContent(iPt)
        eff_val = rec_val / gen_val
        eff_err = np.sqrt(eff_val * (1 - eff_val) / gen_val)
        eff_hist.SetBinContent(iPt, eff_val)
        eff_hist.SetBinError(iPt, eff_err)
    return eff_hist

is_omega = True
rebins = 10
up_to_pt = 7

if is_omega:
    histo_name = "omegaHist"
    outfile_name = "effST_omega.root"
else:
    histo_name = "xiHist"
    outfile_name = "effST_xi.root"

an_file = ROOT.TFile("data/AnalysisResults.root")
histo_full_xi = an_file.Get(f"strangeness-tracking-q-c/{histo_name}")
histo_full_xi.SetDirectory(0)
histo_tracked_xi = an_file.Get(f"strangeness-tracking-q-c/{histo_name}Tracked")
histo_tracked_xi.SetDirectory(0)
## get xy projections
histo_full_xy_xi = histo_full_xi.Project3D("zx")
histo_tracked_xy_xi = histo_tracked_xi.Project3D("zx")
histo_full_xy_xi.RebinX(rebins)
histo_tracked_xy_xi.RebinX(rebins)

# histograms for efficiency
raw_yield_vs_pt_hist = histo_full_xy_xi.ProjectionX("raw_yield_vs_pt")
raw_yield_vs_pt_hist.Reset()
raw_yield_vs_pt_tracked_hist = histo_full_xy_xi.ProjectionX("raw_yield_vs_pt_tracked")
raw_yield_vs_pt_tracked_hist.Reset()


file = ROOT.TFile(outfile_name, "RECREATE")
file.mkdir("cascade_fits")
file.mkdir("tracked_fits")
file.cd()


## loop over pt bins
for ibin in range(1, histo_full_xy_xi.GetNbinsX()+1):
    pt_val = histo_full_xy_xi.GetXaxis().GetBinCenter(ibin)
    if abs(pt_val) < 0.5 or abs(pt_val) > up_to_pt:
        continue

    mass_proj = histo_full_xy_xi.ProjectionY("mass_proj", ibin, ibin)
    mass_proj_tracked = histo_tracked_xy_xi.ProjectionY("mass_proj_tracked", ibin, ibin)

    mass_proj.SetName("mass_proj_pt_{}".format(pt_val))
    mass_proj_tracked.SetName("mass_proj_tracked_pt_{}".format(pt_val))

    ## fit it using roofit
    mass_var = ROOT.RooRealVar("mass", "mass", mass_proj.GetXaxis().GetXmin(), mass_proj.GetXaxis().GetXmax())
    mass_full_hist = ROOT.RooDataHist("mass_full_hist", "mass_full_hist", ROOT.RooArgList(mass_var), mass_proj)
    mass_tracked_hist = ROOT.RooDataHist("mass_tracked_hist", "mass_tracked_hist", ROOT.RooArgList(mass_var), mass_proj_tracked)

    mu_var = ROOT.RooRealVar("mu", "mu", mass_proj.GetXaxis().GetXmin(), mass_proj.GetXaxis().GetXmax())
    frac = ROOT.RooRealVar('f', 'fraction of signal', 0., 1.)
    sigma = ROOT.RooRealVar('sigma', 'casc width', 0.0005, 0.004, 'GeV/c^{2}')
    a1 = ROOT.RooRealVar('a1', 'a1', 0.7, 5.)
    a2 = ROOT.RooRealVar('a2', 'a2', 0.7, 5.)
    n1 = ROOT.RooRealVar('n1', 'n1', 0., 1.)
    n2 = ROOT.RooRealVar('n2', 'n2', 0., 1.)
    c0 = ROOT.RooRealVar('c0', 'constant c0', -100., 100)
    c1 = ROOT.RooRealVar('c1', 'constant c1', -100., 100)
    signal = ROOT.RooCrystalBall('cb', 'cb', mass_var, mu_var, sigma, a1, n1, a2, n2)
    background = ROOT.RooExponential('exp', 'exp', mass_var, c0)

    model = ROOT.RooAddPdf('model', 'model', ROOT.RooArgList(signal, background), ROOT.RooArgList(frac))

    fit_res_full = model.fitTo(mass_full_hist, ROOT.RooFit.Save(True), ROOT.RooFit.PrintLevel(-1))
    file.cd("cascade_fits")
    frame_full = mass_var.frame()
    frame_full.SetName(f"cascade_fit_pt_{pt_val}")
    mass_full_hist.plotOn(frame_full)
    model.plotOn(frame_full)
    model.plotOn(frame_full, ROOT.RooFit.Components("exp"), ROOT.RooFit.LineColor(ROOT.kRed))
    frame_full.Write()
    casc_counts = frac.getVal() * mass_full_hist.sumEntries()
    casc_counts_error = mass_full_hist.sumEntries() * frac.getError()

    fit_res_track = model.fitTo(mass_tracked_hist, ROOT.RooFit.Save(True), ROOT.RooFit.PrintLevel(-1))
    file.cd("tracked_fits")
    frame_tracked = mass_var.frame()
    frame_tracked.SetName(f"tracked_fit_pt_{pt_val}")
    mass_tracked_hist.plotOn(frame_tracked)
    model.plotOn(frame_tracked)
    model.plotOn(frame_tracked, ROOT.RooFit.Components("exp"), ROOT.RooFit.LineColor(ROOT.kRed))
    frame_tracked.Write()
    tracked_counts = frac.getVal() * mass_tracked_hist.sumEntries()
    tracked_counts_error = mass_tracked_hist.sumEntries() * frac.getError()

    print(f"pt: {pt_val}, casc counts: {casc_counts} +- {casc_counts_error}, tracked counts: {tracked_counts} +- {tracked_counts_error}, frac: {frac.getVal()}, error: {frac.getError()}")

    raw_yield_vs_pt_hist.SetBinContent(ibin, casc_counts)
    raw_yield_vs_pt_hist.SetBinError(ibin, casc_counts_error)
    raw_yield_vs_pt_tracked_hist.SetBinContent(ibin, tracked_counts)
    raw_yield_vs_pt_tracked_hist.SetBinError(ibin, tracked_counts_error)

efficiency_hist = computeEfficiency(raw_yield_vs_pt_hist, raw_yield_vs_pt_tracked_hist, "efficiency_hist")
file.mkdir("efficiency")
file.cd("efficiency")
histo_full_xi.Write()
histo_tracked_xi.Write()
histo_full_xy_xi.Write()
histo_tracked_xy_xi.Write()
raw_yield_vs_pt_hist.Write()
raw_yield_vs_pt_tracked_hist.Write()
efficiency_hist.Write()