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)