def plot_equality(title, target, value):
if isinstance(target, torch.Tensor):
target = target.numpy()
if isinstance(value, torch.Tensor):
value = value.numpy()
target = np.squeeze(target)
value = np.squeeze(value)
fig = plt.figure()
plt.stairs(target, baseline=None)
plt.stairs(value, baseline=None)
plt.legend(("Target", "Value"))
plt.title(title)
return fig
def apply_plot(hist, hist_edges):
import sys
import numpy
import matplotlib.pyplot as plt
numpy.set_printoptions(threshold=sys.maxsize)
print("Histogram:")
print(hist)
print("Histogram Edges:")
print(hist_edges)
plt.stairs(hist, hist_edges, fill=True)
plt.xlabel('Tensor value')
plt.ylabel('Counts')
plt.title('Tensor value V.S. Counts')
plt.show()
def plot_inequality(title, value, lower, upper):
if isinstance(value, torch.Tensor):
value = value.numpy()
if isinstance(lower, torch.Tensor):
lower = lower.numpy()
if isinstance(upper, torch.Tensor):
upper = upper.numpy()
value = np.squeeze(value)
lower = np.squeeze(lower)
upper = np.squeeze(upper)
violation = (value > upper) | (lower > value)
fig = plt.figure()
n = np.arange(len(value))
plt.stairs(value, color="b", baseline=None)
plt.stairs(lower, color="r", baseline=None)
plt.stairs(upper, color="g", baseline=None)
plt.scatter(n[violation] + 0.5, value[violation], c="r", marker="|")
plt.legend(("val", "min", "max", "violation"))
plt.title(title)
return fig
}
for h in hists:
hists[h]["bins"] = hists[h]["hist"].GetNbinsX()
hists[h]["contents"] = []
hists[h]["edges"] = []
for i in range(0,hists[h]["bins"]):
w = hists[h]["hist"].GetBinWidth(i)
hists[h]["contents"].append(hists[h]["hist"].GetBinContent(i))
hists[h]["edges"].append(hists[h]["hist"].GetBinCenter(i) - 0.5*w)
#Final upper bin edge
hists[h]["edges"].append(hists[h]["hist"].GetBinCenter(hists[h]["bins"]) + 0.5*w)
#Plot reco MC comparison
fig, ax = plt.subplots(figsize=(8,8))
reco = plt.stairs(hists["h_reco"]["contents"], hists["h_reco"]["edges"], color="crimson", linewidth=2, label="Reco. vertices")
mc = plt.stairs(hists["h_mc"]["contents"], hists["h_mc"]["edges"], color="dodgerblue", linewidth=2, label="MC vertices ($N_{\\mathrm{charged}} > 1$)")
plt.xlim(hists["h_reco"]["xrange"])
plt.ylim(hists["h_reco"]["yrange"])
ax.tick_params(axis='both', which='major', labelsize=25)
plt.xlabel(hists["h_reco"]["xname"],fontsize=30)
plt.ylabel(hists["h_reco"]["yname"],fontsize=30)
l = plticker.MultipleLocator(base=2.0) # this locator puts ticks at regular intervals
ax.xaxis.set_major_locator(l)
plt.legend(fontsize=25)
plt.tight_layout()
fig.savefig(f"{loc.PLOTS}/mc_vs_reco_vertices.pdf")
#Plot the efficiency
fig, ax = plt.subplots(figsize=(8,8))
eff = plt.stairs(hists["h_recoeff_SV_3trk"]["contents"], hists["h_recoeff_SV_3trk"]["edges"], color="k", linewidth=2)
axs[2].set_ylim(-1, 5)
axs[2].set_title("StepPatch artist")
for ax in axs:
ax.legend()
plt.show()
#############################################################################
# *baseline* can take an array to allow for stacked histogram plots
A = [[0, 0, 0],
[1, 2, 3],
[2, 4, 6],
[3, 6, 9]]
for i in range(len(A) - 1):
plt.stairs(A[i+1], baseline=A[i], fill=True)
#############################################################################
# Comparison of `.pyplot.step` and `.pyplot.stairs`
# -------------------------------------------------
#
# `.pyplot.step` defines the positions of the steps as single values. The steps
# extend left/right/both ways from these reference values depending on the
# parameter *where*. The number of *x* and *y* values is the same.
#
# In contrast, `.pyplot.stairs` defines the positions of the steps via their
# bounds *edges*, which is one element longer than the step values.
bins = np.arange(14)
centers = bins[:-1] + np.diff(bins) / 2
y = np.sin(centers / 2)
for ax in axs:
ax.legend()
plt.show()
#############################################################################
# Comparison of `.pyplot.step` and `.pyplot.stairs`.
bins = np.arange(14)
centers = bins[:-1] + np.diff(bins) / 2
y = np.sin(centers / 2)
plt.step(bins[:-1], y, where='post', label='Step(where="post")')
plt.plot(bins[:-1], y, 'o--', color='grey', alpha=0.3)
plt.stairs(y - 1, bins, baseline=None, label='Stairs')
plt.plot(centers, y - 1, 'o--', color='grey', alpha=0.3)
plt.plot(np.repeat(bins, 2),
np.hstack([y[0], np.repeat(y, 2), y[-1]]) - 1,
'o',
color='red',
alpha=0.2)
plt.legend()
plt.title('plt.step vs plt.stairs')
plt.show()
#############################################################################
#
# ------------
#
axs[0].set_title("Step Histograms")
axs[1].stairs(np.arange(1, 6, 1), fill=True,
label='Filled histogram\nw/ automatic edges')
axs[1].stairs(np.arange(1, 6, 1)*0.3, np.arange(2, 8, 1),
orientation='horizontal', hatch='//',
label='Hatched histogram\nw/ horizontal orientation')
axs[1].set_title("Filled histogram")
patch = StepPatch(values=[1, 2, 3, 2, 1],
edges=range(1, 7),
label=('Patch derived underlying object\n'
'with default edge/facecolor behaviour'))
axs[2].add_patch(patch)
axs[2].set_xlim(0, 7)
axs[2].set_ylim(-1, 5)
axs[2].set_title("StepPatch artist")
A = [[0, 0, 0],
[1, 2, 3],
[2, 4, 6],
[3, 6, 9]]
for i in range(len(A) - 1):
plt.stairs(A[i+1], baseline=A[i], fill=True)
plt.show()
# Source: https://matplotlib.org/stable/gallery/lines_bars_and_markers/stairs_demo.html?highlight=stack
def run(nz, ntoys):
#Fetch signal and bkg yields from optimisation
with open(f'{loc.JSON}/optimal_yields_{nz}.json') as f:
yields = json.load(f)
modes = OrderedDict()
file_prefix = "p8_ee_Zbb_ecm91_EvtGen"
#Signal and B+ modes
modes["Bc2TauNu"] = {
"name": f"{file_prefix}_Bc2TauNuTAUHADNU",
"color": "#b2182b",
"label": "$B_c^+ \\to \\tau^+ \\nu_\\tau$",
"N": yields["N_Bc2TauNu"]
}
modes["Bu2TauNu"] = {
"name": f"{file_prefix}_Bu2TauNuTAUHADNU",
"color": "#fdae61",
"label": "$B^+ \\to \\tau^+ \\nu_\\tau$",
"N": yields["N_Bu2TauNu"],
"N_err": yields["N_Bu2TauNu"] * 0.05
} #Belle II B+ -> tau nu precision
modes["Zbb"] = {"name": "p8_ee_Zbb_ecm91"}
#Background decays
for b in bkg_modes:
for d in bkg_modes[b]:
modes[f"{b}_{d}"] = {
"name": f"{file_prefix}_{b}2{d}",
"N": yields[f"N_{b}2{d}"]
}
#Total bkg yield
N_bkg_tot = 0.
for b in bkg_modes:
for d in bkg_modes[b]:
N_bkg_tot += modes[f"{b}_{d}"]["N"]
print("Total bkg expected: %s" % N_bkg_tot)
#Load dataframes for each mode to make templates
df = {}
for m in modes:
df[m] = pd.read_pickle(f"{loc.PKL}/{m}_selected_for_fit.pkl")
#Fit variables
fit_vars = {
"EVT_ThrustEmax_E": {
"name": "Maximum hemisphere E",
"low": 22.,
"high": 52.,
"unit": "GeV/$c^2$"
},
}
#Number of bins in each variable
#30 bins for the 5e12 Z's scenario, scale down for lower lumi
bins = int(30 * np.sqrt(float(nz) / 5.))
#Histogram templates for each mode and variable
#Templates are normalised
h = {}
bin_edges = {}
bin_centres = {}
bin_width = {}
for m in modes:
for v in fit_vars:
h[f"{m}_{v}"], bin_edges[v] = create_hist(
df[m],
bins, [v],
ranges=[[fit_vars[v]["low"], fit_vars[v]["high"]]],
normalise=True,
with_edges=True)
bin_centres[v] = (bin_edges[v][0][1:] + bin_edges[v][0][:-1]) / 2
bin_width[v] = bin_edges[v][0][1] - bin_edges[v][0][0]
#Combine the background histograms according to relative yields from the optimisation
for v in fit_vars:
h[f"bkg_{v}"] = 0
for b in bkg_modes:
for d in bkg_modes[b]:
h[f"bkg_{v}"] += modes[f"{b}_{d}"]["N"] * h[f"{b}_{d}_{v}"]
#Normalise to the total number of bkg events
h[f"bkg_{v}"] = h[f"bkg_{v}"] / N_bkg_tot
#Create a total histgoram of signal + background and then Poisson vary each bin to make a toy dataset
tot_hist = {}
data = {}
data_err = {}
#Make toy datasets
n_toys = int(ntoys)
for i in range(0, n_toys):
np.random.seed(i + 1)
for v in fit_vars:
tot_hist[v] = modes["Bc2TauNu"]["N"] * h[f"Bc2TauNu_{v}"] + modes[
"Bu2TauNu"]["N"] * h[f"Bu2TauNu_{v}"] + N_bkg_tot * h[
f"bkg_{v}"] # h[f"Zbb_{v}"]
data[f"{i}_{v}"] = np.random.poisson(tot_hist[v])
data_err[f"{i}_{v}"] = np.sqrt(data[f"{i}_{v}"])
#Make the total template for the fit in each dimension
def get_template(yield_Bc, yield_Bu, yield_bkg, var):
return yield_Bc * h[f"Bc2TauNu_{var}"] + yield_Bu * h[
f"Bu2TauNu_{var}"] + yield_bkg * h[f"bkg_{var}"]
def binned_nll(template, sample_hist):
return np.sum(template - sample_hist +
sample_hist * np.log((sample_hist + 1e-14) /
(template + 1e-14)))
# 1e-14 added in case there are empty bins
#Loop over toys
results_dict = {}
results_dict["N_Bc"] = {}
results_dict["N_Bu"] = {}
results_dict["N_bg"] = {}
for i in range(0, n_toys):
#Loss function including nll for each of the fit dimensions
def loss(x):
# by default, `x` is an `OrderedSet` of
# zfit parameters.
x = np.array(x)
#print("Value of the parameters", x) # can be commented out, just to see how x evolves during
# the minimisation
# The first parameter is the yield of the Bc signal template
yield_Bc = x[0]
# The second parameter is the yield of the Bu template
yield_Bu = x[1]
# The third parameter is the yield of the bkg template
yield_bkg = x[2]
template = {}
nll = {}
tot_nll = 0
for v in fit_vars:
template[v] = get_template(yield_Bc, yield_Bu, yield_bkg, v)
nll[v] = binned_nll(template[v], data[f"{i}_{v}"])
tot_nll += nll[v]
#Gaussian constraint on B+ -> tau nu yield
tot_nll += (yield_Bu - modes["Bu2TauNu"]["N"]
)**2 / 2. / modes["Bu2TauNu"]["N_err"]**2
return tot_nll
loss.errordef = 0.5 # 0.5 for a log-likelihood, 1 for chi2
#Starting values for the yields
initial_params = {
'value':
[modes["Bc2TauNu"]["N"], modes["Bu2TauNu"]["N"], N_bkg_tot],
'lower': [-1000., -1000., -1000.], # optional
'upper': [100000., 100000., 100000.], # optional
'name': [f'N_Bc_{i}', f'N_Bu_{i}', f'N_bg_{i}'] # optional
}
minimiser = zfit.minimize.Minuit(verbosity=5)
#Since we're using numpy histograms, we need to disable the graph mode of zfit
zfit.run.set_autograd_mode(False)
zfit.run.set_graph_mode(False)
result = minimiser.minimize(loss, initial_params)
param_hesse = result.hesse() # Computation of the errors
corr = result.correlation(method="minuit_hesse")
print(corr)
print(result.info['original'])
params = result.params
print(params)
for p in params:
results_dict["%s" % (p.name[0:4])][f"{i}"] = [
params[p]['value'], param_hesse[p]
]
#Plot first toy
if (n_toys == 1):
for v in fit_vars:
fig, ax = plt.subplots(figsize=(10, 8))
#Plot the toy dataset
Data = plt.errorbar(x=bin_centres[v],
y=data[f"{i}_{v}"],
yerr=data_err[f"{i}_{v}"],
fmt="o",
markersize=4,
color="k") #,label="Generated data")
#Total
h_tot = results_dict["N_Bc"][f"{i}"][0] * h[
f"Bc2TauNu_{v}"] + results_dict["N_Bu"][f"{i}"][0] * h[
f"Bu2TauNu_{v}"] + results_dict["N_bg"][f"{i}"][0] * h[
f"bkg_{v}"]
#h_tot = results_dict["N_Bc"][f"{i}"][0] * h[f"Bc2TauNu_{v}"] + results_dict["N_Bu"][f"{i}"][0] * h[f"Bu2TauNu_{v}"] + results_dict["N_bg"][f"{i}"][0] * h[f"Zbb_{v}"]
Bc = plt.stairs(
h_tot,
bin_edges[v][0],
color=modes["Bc2TauNu"]["color"],
fill=True,
alpha=0.8) #, label=modes["Bc2TauNu"]["label"])
Total = plt.stairs(h_tot,
bin_edges[v][0],
color="k",
linewidth=2) #, label="Total fit")
#Bu
h_Bu = results_dict["N_Bu"][f"{i}"][0] * h[
f"Bu2TauNu_{v}"] + results_dict["N_bg"][f"{i}"][0] * h[
f"bkg_{v}"]
#h_Bu = results_dict["N_Bu"][f"{i}"][0] * h[f"Bu2TauNu_{v}"] + results_dict["N_bg"][f"{i}"][0] * h[f"Zbb_{v}"]
Bu = plt.stairs(
h_Bu,
bin_edges[v][0],
color=modes["Bu2TauNu"]["color"],
fill=True) #, label=modes["Bu2TauNu"]["label"])
#Bkg
h_bkg = results_dict["N_bg"][f"{i}"][0] * h[f"bkg_{v}"]
#h_bkg = results_dict["N_bg"][f"{i}"][0] * h[f"Zbb_{v}"]
Bkg = plt.stairs(
h_bkg, bin_edges[v][0], color="#2166ac", fill=True
) #, label="Background") #label="$Z \\to B^0/B^+/B_s^0/\\Lambda_b^0 X$")
plt.legend(
(Data, Total, Bc, Bu, Bkg),
("Generated data", "Total fit", modes["Bc2TauNu"]["label"],
modes["Bu2TauNu"]["label"], "Background"),
fontsize=22,
loc="upper left")
ax.tick_params(axis='both', which='major', labelsize=25)
plt.xlim(fit_vars[v]["low"], fit_vars[v]["high"])
if (fit_vars[v]["unit"] != ""):
unit_str = "[%s]" % fit_vars[v]["unit"]
unit_space = " "
else:
unit_str = ""
unit_space = ""
plt.xlabel(fit_vars[v]["name"] + " %s" % unit_str, fontsize=30)
plt.ylabel("Candidates / (%.1f%s%s)" %
(bin_width[v], unit_space, fit_vars[v]["unit"]),
fontsize=30)
#plt.yscale('log')
ymin, ymax = plt.ylim()
plt.ylim(0., 1.2 * ymax)
#plt.legend(fontsize=22,loc="upper left")
plt.tight_layout()
fig.savefig(f"{loc.PLOTS}/{v}_template_fit_{nz}.pdf")
#Plot the signal, B+ and background histograms normalised for comparison
fig, ax = plt.subplots(figsize=(10, 8))
h_Bc = h[f"Bc2TauNu_{v}"]
Bc = plt.stairs(h_Bc,
bin_edges[v][0],
color=modes["Bc2TauNu"]["color"],
linewidth=2)
h_Bu = h[f"Bu2TauNu_{v}"]
Bu = plt.stairs(h_Bu,
bin_edges[v][0],
color=modes["Bu2TauNu"]["color"],
linewidth=2)
h_bkg = h[f"bkg_{v}"]
Bkg = plt.stairs(h_bkg,
bin_edges[v][0],
color="#2166ac",
linewidth=2)
#h_inc = h[f"Zbb_{v}"]
#Inc = plt.stairs(h_inc, bin_edges[v][0], color="#92c5de", linewidth=2)
plt.legend(
(Bc, Bu, Bkg), #, Inc),
(modes["Bc2TauNu"]["label"], modes["Bu2TauNu"]["label"],
"Exclusive background"
), #, "Incluisve $Z \\to b\\bar{b}$"),
fontsize=22,
loc="upper left")
ax.tick_params(axis='both', which='major', labelsize=25)
plt.xlim(fit_vars[v]["low"], fit_vars[v]["high"])
if (fit_vars[v]["unit"] != ""):
unit_str = "[%s]" % fit_vars[v]["unit"]
unit_space = " "
else:
unit_str = ""
unit_space = ""
plt.xlabel(fit_vars[v]["name"] + " %s" % unit_str, fontsize=30)
plt.ylabel("Density / (%.1f%s%s)" %
(bin_width[v], unit_space, fit_vars[v]["unit"]),
fontsize=30)
#plt.yscale('log')
ymin, ymax = plt.ylim()
plt.ylim(0., 1.2 * ymax)
#plt.legend(fontsize=22,loc="upper left")
plt.tight_layout()
fig.savefig(f"{loc.PLOTS}/{v}_template_compare_{nz}.pdf")
#Store toy results to json
if (n_toys != 1):
with open(f'{loc.JSON}/toy_template_fit_results_{nz}.json', 'w') as fp:
json.dump(results_dict, fp)
