import ROOT #from ROOT import TTree, TFile, Double import numpy as np from pdg_const import pdg import Tkinter import matplotlib matplotlib.use("agg") import matplotlib.pyplot as plt import pickle as pkl import sys import time from helperfunctions import display_time, prepare_plot import cmath as c import scipy.integrate as integrate from array import array as arr mmu = pdg['muon_M'] mb = pdg["bquark_M"] ms = pdg["squark_M"] mK = pdg["Ks_M"] mB = pdg["Bplus_M"] class model: def __init__(self): #Initialize constants self.mmu = pdg['muon_M'] self.mb = pdg["bquark_M"] self.ms = pdg["squark_M"] self.mK = pdg["Ks_M"] self.mB = pdg["Bplus_M"] self.C7eff = pdg["C7eff"] self.C9eff = pdg["C9eff"] self.C10eff = pdg["C10eff"] self.number_of_decays = pdg["number_of_decays"] self.nr_of_part_list = [] self.nr_of_part_list_names = [] self.resolution = 7.0 #self.C1 = pdg["C1"] #self.C2 = pdg["C2"] #self.C3 = pdg["C3"] #self.C4 = pdg["C4"] self.GF = pdg["GF"] #Fermi coupling const. self.alpha_ew = pdg["alpha_ew"] self.Vts = pdg["Vts"] self.Vtb = pdg["Vtb"] self.param_str = [] self.param_val = [] self.param_val_orig = [] #helperconstants self.x_min = 2*self.mmu self.x_max = (self.mB - self.mK) - 0.1 self.total_pdf_string = "0" self.mode = "" self.param_str.append("total_scale_amp") self.param_val.append(1.0) self.param_str.append("_mean") self.param_val.append(0.0) self.res_counter = 0 self.reco_steps = 10000 #Cusp info self.param_str.append("cusp mean") self.param_val.append(1.0) self.param_str.append("cusp sigma_L") self.param_val.append(1.0) self.param_str.append("cusp sigma_R") self.param_val.append(1.0) self.param_str.append("cusp amp") self.param_val.append(0.0) self.param_val_orig = self.param_val def formfactor(self, q2, subscript): #returns real value #check if subscript is viable if subscript != "0" and subscript != "+" and subscript != "T": raise ValueError('Wrong subscript entered, choose either 0, + or T') #get constants mh = self.mK mbstar0 = pdg["mbstar0"] mbstar = pdg["mbstar"] b0 = pdg["b0"] bplus = pdg["bplus"] bT = pdg["bT"] #N comes from derivation in paper N = 3 #some helperfunctions tpos = (self.mB - self.mK)**2 tzero = (self.mB + self.mK)*(np.sqrt(self.mB)-np.sqrt(self.mK))**2 z_oben = np.sqrt(tpos - q2) - np.sqrt(tpos - tzero) z_unten = np.sqrt(tpos - q2) + np.sqrt(tpos - tzero) z = z_oben/z_unten #calculate f0 if subscript == "0": prefactor = 1/(1 - q2/(mbstar0**2)) _sum = 0 for i in range(N): _sum += b0[i]*(z**i) return prefactor * _sum #calculate f+ or fT else: prefactor = 1/(1 - q2/(mbstar**2)) _sum = 0 if subscript == "T": b = bT else: b = bplus for i in range(N): _sum += b[i] * (z**i - ((-1)**(i-N)) * (i/N) * z**N) return prefactor * _sum def axiv_nonres(self, q2): #returns real value #Load constants GF = self.GF alpha_ew = self.alpha_ew Vtb = self.Vtb Vts = self.Vts C10eff = self.C10eff mmu = self.mmu mb = self.mb ms = self.ms mK = self.mK mB = self.mB #Some helperfunctions beta = np.sqrt(abs(1. - 4. * self.mmu**2. / q2)) kabs = np.sqrt(mB**2 + q2**2/mB**2 + mK**4/mB**2 - 2 * (mB**2 * mK**2 + mK**2 * q2 + mB**2 * q2) / mB**2) #prefactor in front of whole bracket prefactor1 = GF**2. *alpha_ew**2. * (abs(Vtb*Vts))**2 * kabs * beta / (128. * np.pi**5.) #left term in bracket bracket_left = 2./3. * kabs**2 * beta**2 * abs(C10eff*self.formfactor(q2, "+"))**2 #middle term in bracket _top = 4. * mmu**2 * (mB**2 - mK**2) * (mB**2 - mK**2) _under = q2 * mB**2 bracket_middle = _top/_under * abs(C10eff * self.formfactor(q2, "0"))**2 #Note sqrt(q2) comes from derivation as we use q2 and plot q return self.param_val[0]*prefactor1 * (bracket_left + bracket_middle) * 2 * np.sqrt(q2) def vec_nonres(self, q2): #returns real value #Load constants GF = self.GF alpha_ew = self.alpha_ew Vtb = self.Vtb Vts = self.Vts C7eff = self.C7eff C9eff = self.C9eff mmu = self.mmu mb = self.mb ms = self.ms mK = self.mK mB = self.mB #Some helperfunctions beta = np.sqrt(abs(1. - 4. * self.mmu**2. / q2)) kabs = np.sqrt(mB**2 + q2**2/mB**2 + mK**4/mB**2 - 2 * (mB**2 * mK**2 + mK**2 * q2 + mB**2 * q2) / mB**2) #prefactor in front of whole bracket prefactor1 = GF**2. *alpha_ew**2. * (abs(Vtb*Vts))**2 * kabs * beta / (128. * np.pi**5.) #right term in bracket prefactor2 = kabs**2 * (1. - 1./3. * beta**2) abs_bracket = abs(C9eff * self.formfactor(q2, "+") + 2 * C7eff * (mb + ms)/(mB + mK) * self.formfactor(q2, "T"))**2 bracket_right = prefactor2 * abs_bracket #Note sqrt(q2) comes from derivation as we use q2 and plot q return self.param_val[0]*prefactor1 * bracket_right * 2 * np.sqrt(q2) def total_nonres(self, q2, absolut = False): #returns complex value #Load constants GF = self.GF alpha_ew = self.alpha_ew Vtb = self.Vtb Vts = self.Vts C10eff = self.C10eff C9eff = self.C9eff C7eff = self.C7eff mmu = self.mmu mb = self.mb ms = self.ms mK = self.mK mB = self.mB #vector nonresonant part vec_nonres_part = self.vec_nonres(q2) #axial verctor nonresonant part including C7 axiv_nonres_part = self.axiv_nonres(q2) #Complete term if absolut: return vec_nonres_part + axiv_nonres_part else: return complex(vec_nonres_part + axiv_nonres_part, 0.0) def generate_points(self, set_size = 10000, x_min = 2* mmu, x_max = (mB - mK) - 0.1, save = True, verbose = 1, mode = "true_data", resolution = 7.0, nonres_set_size = 44000, nr_of_toys = 1): #returns part_set (=dic) #Check if entered mode and verbose are valid if mode != "slim_points" and mode != "fast_binned" and mode != "true_data": raise ValueError('Invalid mode entered, choose either slim_points, fast_binned or true_data') if verbose != 0 and verbose != 1: raise ValueError('Invalid verbose entered, choose either 0 (no verbose) or 1 (verbose)') #Setup contants and variables self.mode = mode mB = self.mB mK = self.mK mmu = self.mmu #Prepare some variables used in all modes _max = self.total_pdf(3096.0**2) #slim_mode: Generates data of a set size in range(x_min,x_max) according to total_pdf -> Only saves the valid q values (Hit and miss method) if mode == "slim_points": #Prepare variables x_part = [] print("Generating set of size {}...".format(int(set_size))) #Change seed here if necessary # ROOT.TRandom1().SetSeed(0) #Preparecalls for verbose percentage counter if verbose != 0: verbose_calls = [] j = 0 o = 0 for j in range(100): verbose_calls.append(int(set_size*(j+1)/100)) #Get start time for projected time and other variables start = time.time() counter = 0 counter_x = 0 #Loop until set has wanted size while len(x_part) < set_size: counter += 1 x = ROOT.TRandom1().Uniform(x_min, x_max) y = ROOT.TRandom1().Uniform(0, _max) if y < self.total_pdf(x**2): x_part.append(x) counter_x += 1 #progress calls if verbose != 0: end = time.time() if o < 100 and counter%5000 == 0: print(" {0}{1} completed".format(o, "%")) print(" Projected time left: {0}".format(display_time(int((end-start)*set_size/(counter_x+1)-(end-start))))) sys.stdout.write("\033[F") sys.stdout.write("\x1b[2K") sys.stdout.write("\033[F") sys.stdout.write("\x1b[2K") if o >=100: print(" Time to generate set: {0}".format(display_time(int(end-start)))) if len(x_part) == verbose_calls[o]: o += 1 #Generate random points counter += 1 x = ROOT.TRandom1().Uniform(x_min, x_max) y = ROOT.TRandom1().Uniform(0, _max) #Add particles if under pdf if y < self.total_pdf(x**2): x_part.append(x) counter_x += 1 #progress calls if verbose != 0: end = time.time() if o < 100 and counter%5000 == 0: print(" {0}{1} completed".format(o, "%")) print(" Projected time left: {0}".format(display_time(int((end-start)*set_size/(counter_x+1)-(end-start))))) sys.stdout.write("\033[F") sys.stdout.write("\x1b[2K") sys.stdout.write("\033[F") sys.stdout.write("\x1b[2K") if o >=100: print(" Time to generate set: {0}".format(display_time(int(end-start)))) if counter_x == verbose_calls[o]: o += 1 print(" {0} out of {1} particles chosen!".format(int(counter_x), counter)) print(" Set generated!") #Save the set if save: part_set = {"x_part" : x_part} pkl.dump( part_set, open("./data/slim_points/set_{0}_range({1}-{2}).pkl".format(int(set_size),int(x_min), int(x_max)) , "wb" ) ) print(" Set saved!") print #returns all the chosen points (x_part, y_part) and all the random points generated (x, y) return part_set elif mode == "true_data": #Change Seed here if necessary #ROOT.TRandom1().SetSeed(0) #Range of the whole nonres spectrum (used to count total hits in the nonres part) if x_min < 3096.0 and x_max > 3096.0: _max = self.total_pdf(3096.0**2) elif self.total_pdf(x_min**2) < self.total_pdf(x_max**2): _max = self.total_pdf(x_max**2) else: _max = self.total_pdf(x_min**2) x_min_nr = self.x_min x_max_nr = self.x_max _start = time.time() _end = _start tot = [] for toy_nr in range(nr_of_toys): tot.append(0) for toy_nr in range(nr_of_toys): #Prepare variables o = 0 x_part = [] _nonres_set_size = int(ROOT.TRandom1().Gaus(nonres_set_size, np.sqrt(nonres_set_size))) print("Generating toy set {1} of {2} with {0} total rare nonresonant particles...".format(int(_nonres_set_size), toy_nr, nr_of_toys - 1)) #Take starting time and other variables for projected time start = time.time() counter = 0 counter_x = 0 counter_x_nr = 0 #Prepare verbose calls if verbose != 0: verbose_calls = [] o = 0 for j in range(100): verbose_calls.append(int(_nonres_set_size*(j+1)/100)) #Loop until enough particles in nonresonant part while counter_x_nr < _nonres_set_size: #Generate random points x = ROOT.TRandom1().Uniform(x_min_nr, x_max_nr) y = ROOT.TRandom1().Uniform(0, _max) #Add the part to the set if within range(x_min,x_max) and under pdf if x > x_min and x < x_max and y < self.total_pdf(x**2): x_part.append(x) counter_x += 1 #Count the hits in the nonres part -> Range of the whole pdf! if y < self.total_nonres(x**2, absolut = True): counter_x_nr += 1 #Calculate total estimated time to produce all toys end = time.time() tot[toy_nr] = (end-start)*_nonres_set_size/(counter_x_nr+1) total_estim = nr_of_toys*np.mean(tot[:toy_nr+1]) #progress calls if verbose != 0: if o < 100 and counter%10000 == 0: print(" {0}{1} completed of this toy".format(o, "%")) print(" {0} rare nonresonant particles".format(counter_x_nr)) print(" {0} total particles".format(counter_x)) print(" Projected time left for this toy: {0}".format(display_time(int((end-start)*_nonres_set_size/(counter_x_nr+1)-(end-start))))) print(" Projected time left for all toys: {0}".format(display_time(int( total_estim - (end-_start) )))) sys.stdout.write("\033[F") sys.stdout.write("\x1b[2K") sys.stdout.write("\033[F") sys.stdout.write("\x1b[2K") sys.stdout.write("\033[F") sys.stdout.write("\x1b[2K") sys.stdout.write("\033[F") sys.stdout.write("\x1b[2K") sys.stdout.write("\033[F") sys.stdout.write("\x1b[2K") if o >=100: print(" Time to generate this toy set: {0}".format(display_time(int(end-start)))) if counter_x_nr == verbose_calls[o]: o += 1 counter += 1 _end = time.time() print(" {0} out of {1} particles chosen!".format(int(counter_x), counter)) print(" Toy set {0} generated!".format(toy_nr)) #Save the set if save: part_set = {"x_part" : x_part, "counter_x_nr": counter_x_nr, "x_min": x_min, "x_max": x_max, "mode": mode} pkl.dump( part_set, open("./data/true_data/true_data_toy_{0}_range({1}-{2}).pkl".format(toy_nr ,int(x_min), int(x_max)) , "wb" ) ) print(" Toy set {0} saved!".format(toy_nr)) print #returns all the chosen points (x_part, y_part) and all the random points generated (x, y) return part_set elif mode == "fast_binned": #Calculate number of bins nbins = int((x_max-x_min)/resolution) for toy_nr in range(nr_of_toys): print("Generating toy set {1} of {2} with {0} bins...".format(nbins, toy_nr, nr_of_toys)) #Get the mean dq values for each bin and get value of pdf there dq = np.linspace(x_min, x_max ,nbins+1) bin_mean = [] bin_true_height = [] for i in range(len(dq)-1): a = dq[i] b = dq[i+1] bin_mean.append((a+b)/2) _a, _e = integrate.quad(self.total_pdf, a**2, b**2, limit = 200) height = _a/(b-a) bin_true_height.append(height) #Vary the nonres_set_size _nonres_set_size = int(ROOT.TRandom1().Gaus(nonres_set_size, np.sqrt(nonres_set_size))) #Scale true bin size according to nonres_set_size _area, _area_err = integrate.quad(lambda x: self.total_nonres(x ,absolut = True), self.x_min**2, self.x_max**2, limit = 3000) _mean_scan = _area/(self.x_max - self.x_min) _mean_histo_nonres = _nonres_set_size/nbins _ = _mean_histo_nonres/_mean_scan # print(_) for i in range(len(bin_true_height)): bin_true_height[i] = bin_true_height[i]*_ #Start time start = time.time() bin_height = [] #Draw random numbers for i in range(len(bin_mean)): bin_height.append(int(ROOT.TRandom1().Gaus(bin_true_height[i], np.sqrt(bin_true_height[i])))) #print(int(ROOT.TRandom1().Gaus(bin_true_height[i], np.sqrt(bin_true_height[i])))) #progress calls if verbose != 0: end = time.time() print(" Time to generate set: {0}s".format(end-start)) print(" {0} bins simulated".format(nbins)) print(" Set generated!") #Save the set if save: _ = 0 for i in bin_height: _ += i part_set = {"bin_mean" : bin_mean, "bin_height": bin_height, "nbins": nbins, "bin_true_height": bin_true_height, "_nonres_set_size": nonres_set_size, "x_min": x_min, "x_max": x_max} pkl.dump( part_set, open("./data/fast_binned/binned_toy_{0}.pkl".format(toy_nr) , "wb" ) ) print(" Set saved!") print return part_set def draw_plots(self, part_set, mode, x_min = 2* mmu, x_max = (mB - mK) - 0.1, resolution = 7.0): #draws and saves plots, returns nothing #Resolution based on experiment chosen to be ~7MeV #check for invalid entries if mode != "slim_points" and mode != "fast_binned" and mode != "true_data": raise ValueError('Wrong mode entered, choose either slim_points, fast_binned or true_data' ) if mode != self.mode: raise ValueError('self.mode and mode are different, set them to the same value') print("Generating plots") if mode == "fast_binned": #Load variables mB = self.mB mK = self.mK mmu = self.mmu #Range of function in MeV #Choosing the number of steps greatly varies the mean of dgamma!!! 5000 -> 10000 yields a factor of 1.68 dq = np.linspace(x_min, x_max ,self.reco_steps) #Translate to MeV**2 dq2 = [] for i in dq: dq2.append(i**2) #calculate formfactors ff_plus = [] ff_T = [] ff_0 = [] for i in dq: ff_0.append(self.formfactor(i**2, "0")) ff_T.append(self.formfactor(i**2, "T")) ff_plus.append(self.formfactor(i**2, "+")) #calculate nonresonant dgamma_axiv_nonres = [] dgamma_vec_nonres = [] dgamma_total_nonres = [] dgamma_tot = [] for i in dq: dgamma_axiv_nonres.append(self.axiv_nonres(i**2)) dgamma_vec_nonres.append(self.vec_nonres(i**2)) dgamma_total_nonres.append(abs(self.total_nonres(i**2))) dgamma_tot.append(self.total_pdf(i**2)) #Plot formfactors plt.clf() plt.plot(dq2, ff_0, label = "0") plt.plot(dq2, ff_T, label = "T") plt.plot(dq2, ff_plus, label = "+") prepare_plot("Formfactors") plt.savefig("./plots/fast_binned/ff.png") print(" ff.png created") #Plot nonresonant part plt.clf() plt.plot(dq, dgamma_axiv_nonres, label = "axiv") plt.plot(dq, dgamma_vec_nonres, label = "vec") plt.plot(dq, dgamma_total_nonres, label = "total_nonres") prepare_plot("Nonresonant parts") plt.savefig("./plots/fast_binned/vec_axiv.png") print(" vec_axiv.png created") #Plot all pdfs and parts tot_y = [] jpsi_y = [] psi2s_y = [] total_nonres_y = [] cusp_y = [] jpsi_mass, jpsi_width, jpsi_phase, jpsi_scale = pdg["jpsi"] psi2s_mass, psi2s_width, psi2s_phase, psi2s_scale = pdg["psi2s"] for i in range(len(dq)): #print(i**2 - 4*(mmu**2)) tot_y.append(abs(self.total_pdf(dq[i]**2))) jpsi_y.append(abs(self.resonance(dq[i]**2, jpsi_mass, jpsi_width, jpsi_phase, jpsi_scale))) psi2s_y.append(abs(self.resonance(dq[i]**2, psi2s_mass, psi2s_width, psi2s_phase, psi2s_scale))) total_nonres_y.append(abs(self.total_nonres(dq[i]**2))) cusp_y.append(abs(self.bifur_gauss(dq[i]**2, self.param_val[2], self.param_val[5], self.param_val[3], self.param_val[4] ))) #resonance(self, q2, _mass, width, phase, scale): #w[i] = np.sqrt(w[i]) #print(test_y[i]) plt.clf() plt.plot(dq, tot_y, label = "total pdf") plt.plot(dq, jpsi_y, label = "jpsi") plt.plot(dq, psi2s_y, label = "psi2s") plt.plot(dq, total_nonres_y, label = "nonres") plt.plot(dq, cusp_y, label = "cusp") prepare_plot("All pdfs") plt.ylim(0, 2*self.param_val[1]) plt.savefig("./plots/fast_binned/pdf_and_parts.png") print(" pdf_and_parts.png created") #Create histo with pdf dq2 = [] for i in dq: dq2.append(i**2) dgamma_tot = [] for i in dq2: dgamma_tot.append(self.total_pdf(i)) #Load particle set bin_mean = part_set["bin_mean"] bin_height = part_set["bin_height"] #Scale the pdf to the histo _sum = np.sum(bin_height) _mean = np.mean(bin_height) # nbins = part_set["nbins"] # _ = self.param_val[0] area, err = integrate.quad(self.total_pdf, x_min**2, x_max**2, limit = 300) _m = area/(x_max-x_min) self.param_val[0] = self.param_val[0]*_mean/_m plt.clf() plt.ylim(0.0, 2.5*_mean) #Choose range in between the 2 resonances # plt.xlim(3150, 3650) #Draw the plot plt.hist(bin_mean, bins=nbins, range=(x_min, x_max), weights = bin_height, label = "toy data binned") plt.plot(dq, dgamma_tot, label = "pdf") prepare_plot("{0} random points generated ({1} particles)".format(len(bin_mean), int(_sum))) plt.savefig("./plots/fast_binned/histo.png") self.param_val[0] = _ print(" histo.png created") print(" All plots drawn \n") return else: #Range of function in MeV dq = np.linspace(x_min, x_max ,self.reco_steps) #Translate to MeV**2 dq2 = [] for i in dq: dq2.append(i**2) #calculate formfactors ff_plus = [] ff_T = [] ff_0 = [] for i in dq: ff_0.append(self.formfactor(i**2, "0")) ff_T.append(self.formfactor(i**2, "T")) ff_plus.append(self.formfactor(i**2, "+")) #calculate nonresonant dgamma_axiv_nonres = [] dgamma_vec_nonres = [] dgamma_total_nonres = [] dgamma_tot = [] for i in dq: dgamma_axiv_nonres.append(self.axiv_nonres(i**2)) dgamma_vec_nonres.append(self.vec_nonres(i**2)) dgamma_total_nonres.append(abs(self.total_nonres(i**2))) dgamma_tot.append(self.total_pdf(i**2)) #Plot formfactors plt.clf() plt.plot(dq2, ff_0, label = "0") plt.plot(dq2, ff_T, label = "T") plt.plot(dq2, ff_plus, label = "+") prepare_plot("Formfactors") plt.savefig("./plots/points/ff.png") print(" ff.png created") #Plot nonresonant part plt.clf() plt.plot(dq, dgamma_axiv_nonres, label = "axiv") plt.plot(dq, dgamma_vec_nonres, label = "vec") plt.plot(dq, dgamma_total_nonres, label = "total_nonres") prepare_plot("Nonresonant parts") plt.savefig("./plots/points/vec_axiv.png") print(" vec_axiv.png created") #Particle set x_part = part_set["x_part"] set_size = len(x_part) #Histo unnormalized bins = int((x_max-x_min)/resolution) plt.clf() _y, _x, _ = plt.hist(x_part, bins=bins, range=(x_min, x_max), label = "toy data binned ({0} points)".format(set_size)) _mean_histo = float(np.mean(_y)) plt.ylim(0.0, 2.5*_mean_histo) prepare_plot("Binned toy data") plt.savefig("./plots/points/histo_raw.png") print(" histo_raw.png created") #Histo and pdf normailzed plt.clf() #Plot histo #Scale the pdf to the histo _ = [] __ = self.param_val[0] for i in range(len(_y)): _.append(x_min+(x_max-x_min)/bins*i) # area, err = integrate.quad(lambda x:self.total_nonres(x, absolut = True), self.x_min**2, self.x_max**2, limit = 300) # _mean = area/(x_max-x_min) # self.param_val[0] = self.param_val[0]*44000.0/_mean self.param_val[0] *= self.number_of_decays _dgamma_tot = [] for i in dq: _dgamma_tot.append(self.total_pdf(i**2)) plt.ylim(0.0, 2.5*_mean_histo) plt.hist(_, bins=bins, range=(x_min, x_max), weights = _y, label = "toy data binned") plt.plot(dq, _dgamma_tot, label = "pdf scaled up to match histo") self.param_val[0] = __ #Only show from 0 to twice the mean of pdf # plt.ylim(0, 2*self.param_val[1]) prepare_plot("{0} random points generated according to pdf".format(set_size)) plt.savefig("./plots/points/histo.png") print(" histo.png created") #Plot all pdfs tot_y = [] jpsi_y = [] psi2s_y = [] total_nonres_y = [] cusp_y = [] jpsi_mass, jpsi_width, jpsi_phase, jpsi_scale = pdg["jpsi"] psi2s_mass, psi2s_width, psi2s_phase, psi2s_scale = pdg["psi2s"] for i in range(len(dq)): #print(i**2 - 4*(mmu**2)) tot_y.append(abs(self.total_pdf(dq[i]**2))) jpsi_y.append(abs(self.resonance(dq[i]**2, jpsi_mass, jpsi_width, jpsi_phase, jpsi_scale))) psi2s_y.append(abs(self.resonance(dq[i]**2, psi2s_mass, psi2s_width, psi2s_phase, psi2s_scale))) total_nonres_y.append(abs(self.total_nonres(dq[i]**2))) cusp_y.append(abs(self.bifur_gauss(dq[i]**2, self.param_val[2], self.param_val[5], self.param_val[3], self.param_val[4] ))) #resonance(self, q2, _mass, width, phase, scale): #w[i] = np.sqrt(w[i]) #print(test_y[i]) plt.clf() # plt.ylim(0,2*_max) plt.plot(dq, tot_y, label = "total pdf") plt.plot(dq, jpsi_y, label = "jpsi") plt.plot(dq, psi2s_y, label = "psi2s") plt.plot(dq, total_nonres_y, label = "nonres") plt.plot(dq, cusp_y, label = "cusp") prepare_plot("All pdfs and parts") plt.savefig("./plots/points/pdf_and_parts.png") print(" pdf_and_parts.png created") print(" All plots drawn \n") return def total_pdf(self, q2): #return absolut value of sum of all complex values (real number) #Calculate the pdf with the added resonances exec("_sum = abs({0})".format(self.total_pdf_string)) return _sum def resonance(self, q2, _mass, width, phase, scale, absolut = False): #returns complex value #Helper variables p = 0.5 * np.sqrt(q2 - 4*(mmu**2)) p0 = 0.5 * np.sqrt(_mass**2 - 4*mmu**2) gamma_j = p / p0 * _mass /q2 * width #Calculate the resonance _top = complex(_mass * width, 0.0) _bottom = complex((_mass**2 - q2), -_mass*gamma_j) com = _top/_bottom * scale #Rotate by the phase r = abs(com) _phase = c.phase(com) _phase += phase x = c.cos(phase)*r y = c.sin(phase)*r com = complex(x,y) if absolut: return abs(self.param_val[0]*com) else: return self.param_val[0]*com def add_resonance(self, _mass, width, phase, scale, name): #adds string to total_pdf and adds constant places for fit #Count the number of resonances added to the pdf -> used to keep track of the parameters _ = self.res_counter #Adds the resonace to the pdf in form of a string (to be executed later) self.param_str.append(name + " mass") self.param_val.append(_mass) self.param_str.append(name + " width") self.param_val.append(width) self.param_str.append(name + " phase") self.param_val.append(phase) self.param_str.append(name + " scale") self.param_val.append(scale) self.param_val_orig = self.param_val self.total_pdf_string += "+ self.resonance(q2,{0},{1},{2},{3})".format(self.param_val[6+_*4], self.param_val[7+_*4],self.param_val[8+_*4], self.param_val[9+_*4]) self.res_counter += 1 def bifur_gauss(self, q2, mean, amp, sigma_L, sigma_R, absolut = False): #returns complex value #Get q out of q2 q = np.sqrt(q2) #Choose the right sigma depending on side of the cusp if q < mean: sigma = sigma_L else: sigma = sigma_R #Calculate the exponential part of the cusp _exp = np.exp(- (q-mean)**2 / (2 * sigma**2)) #Scale so the total area under curve is 1 and the top of the cusp is continuous dgamma = amp*_exp/(np.sqrt(2*np.pi))*2*(sigma_L*sigma_R)/(sigma_L+sigma_R) com = complex(dgamma, 0) if absolut: return abs(self.param_val[0]*com) else: return self.param_val[0]*com def add_cusp(self, mean, amp, sigma_L, sigma_R): #adds string to total_pdf and adds constant places for fit #Add cusp to total_pdf self.total_pdf_string += "+ self.bifur_gauss(q2,{0},{1},{2},{3})".format(mean, "self.param_val[5]", "self.param_val[3]", "self.param_val[4]") #Save the variables of the cusp -> Used for fitting self.param_val[2] = mean self.param_val[3] = sigma_L self.param_val[4] = sigma_R self.param_val[5] = amp self.param_val_orig = self.param_val def normalize_pdf(self, value = 1.0, verbose = 1): #Normalizes pdf with a global factor in front, saves mean and global factor if verbose == 1: print("Normalizing pdf...") area, err = integrate.quad(self.total_pdf, self.x_min**2, self.x_max**2, limit = 3000) _mean = area/(self.x_max-self.x_min) self.param_val[0] = self.param_val[0]/area*value self.param_val[1] = _mean/area*value self.param_val_orig = self.param_val if verbose == 1: print(" pdf normalized!") print def param_list(self, orig = True): #prints a list of all parameters of the form: number, name, value if orig: print("List of parameters (true values)") else: print("List of parameters (fitted values)") print(" Nr. Name Value") for i in range(len(self.param_val_orig)): if orig: print(" {0}. {1}: {2}".format(i, self.param_str[i], self.param_val_orig[i])) else: print(" {0}. {1}: {2}".format(i, self.param_str[i], self.param_val[i])) print def add_nonres(self): self.total_pdf_string += " + self.total_nonres(q2)" def log_likelihood(self, npar, apar, part_set, fitting_param_index, resolution = 7.0): #Replaced soon with TMinuit # self.param_val = apar for i in range(len(fitting_param_index)): self.param_val[fitting_param_index[i]] = apar[i] self.normalize_pdf(verbose = 0) # part_set = {"bin_height": bin_height, "bin_mean": bin_mean, "_x": _x, "nbins": nbins, "x_max": x_max, "x_min": x_min} _x = part_set["_x"] bin_height = part_set["bin_height"] ll = 0.0 for i in range(len(bin_height)): area, _ = integrate.quad(self.total_pdf, _x[i]**2, _x[i+1]**2, limit = 300) y_true = area/(_x[i+1]-_x[i]) ll += np.log(bin_height[i]*y_true) ll = -ll return ll def fit_pdf_to_data(self, part_set, fitting_param_index = [2,3,4,5,9,13], unbinned_data = True, resolution = 0): # x_part = part_set["x_part"] if resolution == 0: resolution = self.resolution if unbinned_data: #Load data x_part = part_set["x_part"] x_min = part_set["x_min"] x_max = part_set["x_max"] #Fill data into bins nbins = int((x_max-x_min)/resolution) bin_height , _x, _ = plt.hist(x_part, bins=nbins, range=(x_min, x_max)) plt.clf() bin_mean = [] for i in range(len(_x)-1): bin_mean.append((_x[i]+_x[i+1])/2) _part_set = {"bin_height": bin_height, "bin_mean": bin_mean, "_x": _x, "nbins": nbins, "x_max": x_max, "x_min": x_min} x_part = arr("f", x_part) nPoints = len(x_part) name = self.param_str npar = len(fitting_param_index) vstart = arr("d", npar*[0]) # print(len(vstart)) amp_err = self.param_val[0]/20 cusp_m_err = 10 step = arr( 'd', ( amp_err, cusp_m_err, 1.0, 0.1, 1e-7, self.param_val[9]/50, self.param_val[13] )) npar = len(vstart) def fcn(npar, deriv, f, apar, iflag): f[0] = self.log_likelihood(npar, apar, _part_set, resolution = 7.0, fitting_param_index = fitting_param_index) # return f[0] gMinuit = ROOT.TMinuit(npar) gMinuit.SetFCN( fcn ) gMinuit.SetErrorDef(0.5) arglist = arr( 'd', npar*[0.] ) ierflg = ROOT.Long(1998) arglist[0] = 1 gMinuit.mnexcm( "SET ERR", arglist, 1, ierflg ) _ = 0 for j in fitting_param_index: # print(_, fitting_param_index[_]) vstart[_] = self.param_val[fitting_param_index[_]] if _ == 1: vstart[_] -= 50 gMinuit.mnparm( _, self.param_str[fitting_param_index[_]], vstart[_], step[_], 0, 0, ierflg ) _ += 1 arglist[0] = 6000 arglist[1] = 1. gMinuit.mnexcm( "MIGRAD", arglist, 2, ierflg ) amin, edm, errdef = 0.18, 0.19, 0.20 nvpar, nparx, icstat = 1983, 1984, 1985 # gMinuit.mnstat( amin, edm, errdef, nvpar, nparx, icstat ) # gMinuit.mnprin( 3, amin ) #__________________________________________________________________________________ # # gMinuit = ROOT.TMinuit(13) # gMinuit.SetFCN( self.log_likelihood(part_set = part_set, ) ) # arglist = arr( 'd', 15*[0.] ) # gMinuit.mnexcm( "SET ERR", arglist, 1, ierflg ) # # ierflg = 0 # # arglist = arr( 'd', (6000, 0.3) ) # # amp_err = np.sqrt(self.param_val[0]) # cusp_m_err = np.sqrt(self.param_val[2]) # # vstart = arr( 'd', self.param_val ) # step = arr( 'd', ( amp_err, 0, cusp_m_err, 1.0, 0.1, 1e-7, 0, 0, 0, self.param_val[9]/50, 0, 0, 0, self.param_val[13] )) # # for i in range(len(self.param_val)): # gMinuit.mnparm( i, "a{0}".format(i), vstart[i], step[i], 0, 0, ierflg ) # gMinuit.mnexcm( "MIGRAD", arglist, 2, ierflg ) # # #__________________________________________________________________________________ # # Error = 0; # z = array( 'f', ( 1., 0.96, 0.89, 0.85, 0.78 ) ) # errorz = array( 'f', 5*[0.01] ) # # x = array( 'f', ( 1.5751, 1.5825, 1.6069, 1.6339, 1.6706 ) ) # y = array( 'f', ( 1.0642, 0.97685, 1.13168, 1.128654, 1.44016 ) ) # # ncount = 0 # # ##______________________________________________________________________________ # def testfit(): # # # # gMinuit = TMinuit(5) # # gMinuit.SetFCN( fcn ) # # # # arglist = array( 'd', 10*[0.] ) # # ierflg = 1982 # # # # arglist[0] = 1 # # gMinuit.mnexcm( "SET ERR", arglist, 1, ierflg ) # # # Set starting values and step sizes for parameters # vstart = array( 'd', ( 3, 1, 0.1, 0.01 ) ) # step = array( 'd', ( 0.1, 0.1, 0.01, 0.001 ) ) # gMinuit.mnparm( 0, "a1", vstart[0], step[0], 0, 0, ierflg ) # gMinuit.mnparm( 1, "a2", vstart[1], step[1], 0, 0, ierflg ) # gMinuit.mnparm( 2, "a3", vstart[2], step[2], 0, 0, ierflg ) # gMinuit.mnparm( 3, "a4", vstart[3], step[3], 0, 0, ierflg ) # # # Now ready for minimization step # arglist[0] = 500 # arglist[1] = 1. # gMinuit.mnexcm( "MIGRAD", arglist, 2, ierflg ) # # # Print results # amin, edm, errdef = 0.18, 0.19, 0.20 # nvpar, nparx, icstat = 1983, 1984, 1985 # gMinuit.mnstat( amin, edm, errdef, nvpar, nparx, icstat ) # gMinuit.mnprin( 3, amin ) # # # ##______________________________________________________________________________ # # def fcn( npar, gin, f, par, iflag ): # global ncount # # nbins = 5 # # # # # calculate chisquare # # chisq, delta = 0., 0. # # for i in range(nbins): # # delta = (z[i]-func(x[i],y[i],par))/errorz[i] # # chisq += delta*delta # # # # f[0] = chisq # ncount += 1 # # # # def func( x, y, par ): # # value = ( (par[0]*par[0])/(x*x)-1)/ ( par[1]+par[2]*y-par[3]*y*y) # # return value # # # ##______________________________________________________________________________ # if __name__ == '__main__': # testfit()