- import ROOT
- #from ROOT import TTree, TFile, Double
- import numpy as np
- from pdg_const import pdg
- import matplotlib
- matplotlib.use("Qt5Agg")
- import matplotlib.pyplot as plt
- import pickle as pkl
- import sys
- import time
- from helperfunctions import display_time
- import cmath as c
-
- 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):
-
-
- 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.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.x_min = 2*self.mmu
- self.x_max = (self.mB - self.mK) - 0.1
- self.total_pdf_string = "self.total_nonres(q2)"
- self.mode = ""
- self.total_scale_amp = 1.0
- self._mean = 0.0
-
- self.cusp_mean = 1
- self.cusp_sigma_L = 1
- self.cusp_sigma_R = 1
- self.cusp_amp = 0
-
-
- def formfactor(self, q2, subscript):
-
- #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 = 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):
-
- 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
-
- return prefactor1 * (bracket_left + bracket_middle) * 2 * np.sqrt(q2)
-
-
- def vec_nonres(self, q2):
-
- 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
-
- return prefactor1 * bracket_right * 2 * np.sqrt(q2)
-
- def total_nonres(self, q2):
-
- #Get 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
-
- return self.total_scale_amp*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 = "slim_points", resolution = 7.0, min_bin_scaling = 100):
-
- #Setup contants and variables
-
- if mode != "slim_points" and mode != "full_points" and mode != "fast_binned":
- raise ValueError('Wrong mode entered, choose either slim_points, full_points or fast_binned')
-
- self.mode = mode
-
- mB = self.mB
- mK = self.mK
- mmu = self.mmu
-
- #Range of function in MeV
-
- dq = np.linspace(x_min, x_max ,5000)
-
- x1 = 2500
- y1 = self.total_pdf(x1**2)
-
- x2 = 4000
- y2 = self.total_pdf(x2**2)
-
- #Translate to MeV**2
-
- dgamma_tot = []
-
- for i in dq:
- dgamma_tot.append(self.total_pdf(i**2))
-
- dq2 = []
-
- for i in dq:
- dq2.append(i**2)
-
- #Generate random points
-
- psi2s_mass, psi2s_width, psi2s_phase, psi2s_scale = pdg["psi2s"]
-
- _max = max(dgamma_tot)
-
- A1_x1 = (_max-y1)*x1
- A23_x1 = y1*x1
-
- A1_x2 = (_max-y2)*x2
- A23_x2 = y2*x2
-
- if mode == "slim_points":
-
- x_part = []
- y_part = []
-
- print("Generating set of size {}...".format(int(set_size)))
-
- #print(len(y))
-
- #ROOT.TRandom1().SetSeed(0)
-
- if verbose != 0:
- verbose_calls = []
- j = 0
- o = 0
- while j < 100:
- j += verbose
- verbose_calls.append(int(set_size*j/100))
-
- start = time.time()
- counter = 0
- counter_x = 0
- 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*verbose+verbose < 100 and counter%300 == 0:
- print(" {0}{1} completed".format(o*verbose+verbose, "%"))
- print(" Projected time left: {0}".format(display_time(int((end-start)*set_size/(len(x_part)+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*verbose + verbose >=100:
- sys.stdout.write("\033[F")
- sys.stdout.write("\x1b[2K")
- print(" Time to generate set: {0}".format(display_time(int(end-start))))
-
- if len(x_part) == 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, "y_part": y_part, "counter_tot": counter, "counter_x": counter_x}
-
- pkl.dump( part_set, open("./data/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 x_part, y_part, counter
-
- if mode == "full_points":
-
- x = []
- y = []
-
- x_part = []
- y_part = []
-
- print("Generating set of size {}...".format(int(set_size)))
-
- #print(len(y))
-
- #ROOT.TRandom1().SetSeed(0)
-
- if verbose != 0:
- verbose_calls = []
- j = 0
- o = 0
- while j < 100:
- j += verbose
- verbose_calls.append(int(set_size*j/100))
-
- start = time.time()
-
- counter = 0
- counter_x = 0
- while len(x_part) < set_size:
- counter += 1
- x.append(ROOT.TRandom1().Uniform(x_min, x_max))
- y.append(ROOT.TRandom1().Uniform(0, _max))
-
- if y[-1] < self.total_pdf(x[-1]**2):
- x_part.append(x)
- y_part.append(y)
- counter_x += 1
-
- #progress calls
- if verbose != 0:
- end = time.time()
- if o*verbose+verbose < 100 and counter%300 == 0:
- print(" {0}{1} completed".format(o*verbose+verbose, "%"))
- print(" Projected time left: {0}".format(display_time(int((end-start)*set_size/(len(x_part)+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*verbose + verbose >=100:
- sys.stdout.write("\033[F")
- sys.stdout.write("\x1b[2K")
- print(" Time to generate set: {0}".format(display_time(int(end-start))))
-
- if len(x_part) == verbose_calls[o]:
- o += 1
-
-
-
- print(" {0} out of {1} particles chosen!".format(len(x_part), counter))
-
- print(" Set generated!")
-
- #Save the set
-
- if save:
-
- part_set = {"x_part" : x_part, "y_part": y_part, "counter": counter}
-
- raw_set = {"x" : x, "y": y, "counter": counter}
-
- pkl.dump( part_set, open("./data/set_{0}.pkl".format(int(set_size)) , "wb" ) )
-
- pkl.dump( part_set, open("./data/set_raw_{0}.pkl".format(int(set_size)) , "wb" ) )
-
- print(" Sets saved!")
-
- print
-
- #returns all the chosen points (x_part, y_part) and all the random points generated (x, y)
-
- return x_part, y_part, counter
-
-
- if mode == "fast_binned":
-
- nbins = int((x_max-x_min)/resolution)
-
- print("Generating set with {} bins...".format(nbins))
-
- 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]
- c = (a+b)/2
- bin_mean.append(c)
-
- height = self.total_pdf(c**2)
- bin_true_height.append(height)
-
- _min = min(bin_true_height)
-
- for i in range(len(bin_true_height)):
- bin_true_height[i] = bin_true_height[i]/_min*min_bin_scaling
-
- start = time.time()
-
- bin_height = []
-
- 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}
-
- pkl.dump( part_set, open("./data/binned_set_{0}bins_{1}part.pkl".format(nbins, _) , "wb" ) )
-
- print(" Set saved!")
-
- print
-
- return bin_mean, bin_height, nbins
-
-
-
-
- def draw_plots(self, part_set, counter, mode,min_bin_scaling = 100, x_min = 2* mmu, x_max = (mB - mK) - 0.1, resolution = 7):
-
- if mode != "slim_points" and mode != "full_points" and mode != "fast_binned":
- raise ValueError('Wrong mode entered, choose either slim_points, full_points or fast_binned')
- if mode != self.mode:
- raise ValueError('self.mode and mode are different, set them to the same value')
- #Resolution based on experiment chosen to be ~7MeV
-
- #Setup contants and variables
-
- print("Generating plots")
-
- if mode == "fast_binned":
-
- mB = self.mB
- mK = self.mK
- mmu = self.mmu
-
- #Range of function in MeV
-
- dq = np.linspace(x_min, x_max ,5000)
-
- #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_tot = []
-
- for i in dq:
- dgamma_axiv_nonres.append(self.axiv_nonres(i**2))
- dgamma_vec_nonres.append(self.vec_nonres(i**2))
- dgamma_tot.append(self.total_pdf(i**2))
-
-
- #Plot formfactors
-
- plt.plot(dq2, ff_0, label = "0")
- plt.plot(dq2, ff_T, label = "T")
- plt.plot(dq2, ff_plus, label = "+")
-
- plt.grid()
-
- plt.title("Formfactors")
-
- plt.legend()
-
- plt.savefig("./plots/fast_binned/ff.png")
-
- print(" ff.png created")
-
- plt.clf()
-
-
- #Plot nonresonant part
-
- plt.plot(dq, dgamma_axiv_nonres, label = "axiv")
- plt.plot(dq, dgamma_vec_nonres, label = "vec")
-
- plt.grid()
-
- plt.title("Nonresonant axial vector and vector parts")
-
- plt.legend()
-
- plt.savefig("./plots/fast_binned/vec_axiv.png")
-
- print(" vec_axiv.png created")
-
- plt.clf()
-
- plt.plot(dq, dgamma_tot, label = "total")
-
- plt.grid()
-
- plt.title("Total pdf")
-
- plt.legend()
-
- plt.savefig("./plots/fast_binned/tot.png")
-
- print(" tot.png created")
-
- #All pdfs
-
- #print(test_x[1]**2 - self.x_min**2)
-
- 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.cusp_mean, self.cusp_amp, self.cusp_sigma_L, self.cusp_sigma_R )))
- #resonance(self, q2, _mass, width, phase, scale):
- #w[i] = np.sqrt(w[i])
- #print(test_y[i])
-
- plt.clf()
-
- plt.title("All pdfs")
-
- #plt.yscale("log")
-
- plt.ylim(0, 1e-5)
-
- plt.grid()
-
- 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")
-
- plt.legend()
-
- plt.savefig("./plots/fast_binned/pdf_and_parts.png")
-
- print(" pdf_and_parts.png created")
-
- #Create histo with pdf
-
- #Translate to MeV**2
-
- plt.clf()
-
- dq2 = []
-
- for i in dq:
- dq2.append(i**2)
-
- dgamma_tot = []
-
- for i in dq2:
- dgamma_tot.append(self.total_pdf(i))
-
- _min = min(dgamma_tot)
-
- for i in range(len(dgamma_tot)):
- dgamma_tot[i] = dgamma_tot[i]/_min*min_bin_scaling
-
- bin_mean, bin_height = part_set
-
- nbins = counter
-
- plt.hist(bin_mean, bins=nbins, range=(self.x_min, self.x_max), weights = bin_height, label = "toy data binned")
-
- plt.plot(dq, dgamma_tot, label = "pdf")
-
- _sum = 0
-
- for i in bin_height:
- _sum += i
-
- #print(_max)
-
- plt.grid()
-
- plt.ylim(0, 2000)
-
- plt.legend()
-
- plt.title("{0} random points generated according to pdf ({1} particles)".format(len(bin_mean), _sum))
-
- plt.savefig("./plots/fast_binned/histo.png")
-
- print(" histo.png created")
-
- print(" All plots drawn \n")
-
- return
-
- else:
-
- mB = self.mB
- mK = self.mK
- mmu = self.mmu
-
- #Range of function in MeV
-
- dq = np.linspace(x_min, x_max ,5000)
-
- #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_tot = []
-
- for i in dq:
- dgamma_axiv_nonres.append(self.axiv_nonres(i**2))
- dgamma_vec_nonres.append(self.vec_nonres(i**2))
- dgamma_tot.append(self.total_pdf(i**2))
-
-
- #Plot formfactors
-
- plt.plot(dq2, ff_0, label = "0")
- plt.plot(dq2, ff_T, label = "T")
- plt.plot(dq2, ff_plus, label = "+")
-
- plt.grid()
-
- plt.title("Formfactors")
-
- plt.legend()
-
- plt.savefig("./plots/points/ff.png")
-
- print(" ff.png created")
-
- plt.clf()
-
-
- #Plot nonresonant part
-
- plt.plot(dq, dgamma_axiv_nonres, label = "axiv")
- plt.plot(dq, dgamma_vec_nonres, label = "vec")
-
- plt.grid()
-
- plt.title("Nonresonant axial vector and vector parts")
-
- plt.legend()
-
- plt.savefig("./plots/points/vec_axiv.png")
-
- print(" vec_axiv.png created")
-
- plt.clf()
-
- plt.plot(dq, dgamma_tot, label = "total")
-
- plt.grid()
-
- plt.title("Total pdf")
-
- plt.legend()
-
- plt.savefig("./plots/points/tot.png")
-
- print(" tot.png created")
-
-
- #Particle set
-
- x_part, y_part = part_set
-
- set_size = len(x_part)
-
- #Plot generated generate_points
-
- #plt.clf()
-
- #plt.plot(x_part, y_part, label = "Random points generated", marker = ".", linewidth = 0)
-
- #plt.plot(dq, dgamma_tot, label = "pdf")
-
- #plt.grid()
-
- #plt.title("Random points generated and pdf")
-
- #plt.legend()
-
- #plt.savefig("./plots/points/points_raw.png")
-
- #print(" points_raw.png created")
-
- #Histo unnormalized
-
- bins = int((x_max-x_min)/resolution)
-
- plt.clf()
-
- wheights = np.ones_like(x_part)
-
- _max = max(dgamma_tot)
-
- x1 = 2500
- y1 = self.total_pdf(x1**2)
- x2 = 4000
- y2 = self.total_pdf(x2**2)
-
- for i in range(len(wheights)):
- if x_part[i] < x1:
- wheights[i] = x1*y1/(x1*_max)
- elif x_part[i] > x2:
- wheights[i] = x2*y2/(x2*_max)
-
- _y, _x, _ = plt.hist(x_part, bins=bins, weights = wheights, range=(x_min, x_max), label = "toy data binned ({0} points)".format(sum(wheights)))
-
- _mean_histo = float(np.mean(_y))
-
- plt.legend()
-
- plt.title("Binned toy data")
-
- plt.savefig("./plots/points/histo_raw.png")
-
- print(" histo_raw.png created")
-
-
- #Histo and pdf normailzed
-
- plt.clf()
-
- for i in range(len(dgamma_tot)):
- dgamma_tot[i] = dgamma_tot[i]/(float(set_size)*_max * 2.0 * mmu / counter)
-
- _mean = np.mean(dgamma_tot)
-
- #Attempt for marked field of std-dev
-
- #dgamma_min = []
- #dgamma_plu = []
- #for i in range(len(dgamma_tot)):
- #dgamma_min.append(dgamma_tot[i]-np.sqrt(dgamma_tot[i]))
- #dgamma_plu.append(dgamma_tot[i]+np.sqrt(dgamma_tot[i]))
-
- #plt.plot(dq, dgamma_min, alpha = 0.5)
-
- #plt.plot(dq, dgamma_plu, alpha = 0.5)
-
- #plt.fill_between(dq, dgamma_min, dgamma_plu)
-
- #Plot histo
-
- plt.hist(x_part, bins=bins, range=(x_min, x_max), weights = wheights/(_mean_histo/_mean), label = "toy data binned")
-
- plt.plot(dq, dgamma_tot, label = "pdf")
-
- #print(_max)
-
- plt.grid()
-
- plt.legend()
-
- plt.ylim(0, 1e-5)
-
- plt.title("{0} random points generated according to pdf".format(sum(wheights)))
-
- plt.savefig("./plots/points/histo.png")
-
- print(" histo.png created")
-
- #All pdfs
-
- #print(test_x[1]**2 - self.x_min**2)
-
- 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.cusp_mean, self.cusp_amp, self.cusp_sigma_L, self.cusp_sigma_R )))
- #resonance(self, q2, _mass, width, phase, scale):
- #w[i] = np.sqrt(w[i])
- #print(test_y[i])
-
- plt.clf()
-
- plt.title("All pdfs")
-
- #plt.yscale("log")
-
- plt.ylim(0, 2*self._mean)
-
- plt.grid()
-
- 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")
-
- plt.legend()
-
- 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, cusp_amp = -1.0, scan = False):
-
- if scan:
- self.normalize_pdf()
-
- if cusp_amp == -1.0:
- cusp_amp = cusp_amp
- else:
- cusp_amp = self.cusp_amp
-
- #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): #returns [real, imaginary]
-
- #calculate the resonance ---------------------------------------------> Formula correct?
-
- #if abs(np.sqrt(q2) - _mass) < 300:
- #return 0., 0.
-
- 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)
-
- #print(q2)
-
- #Teiler erweitert mit kompl. konj.
-
- #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
-
- #_top_im = - _mass**2 * width * gamma_j
-
- #_top_re = _mass * width * (_mass**2 - q2)
-
- #_bottom = (_mass**2 - q2)**2 + _mass**2 * gamma_j**2
-
- #real = _top_re/_bottom
-
- #imaginary = _top_im/_bottom
-
- #com = complex(real, imaginary) * scale
-
- #r = abs(com)
-
- #_phase = c.phase(com)
-
- #_phase += phase
-
- #x = c.cos(phase)*r
- #y = c.sin(phase)*r
-
- #com = complex(x,y)
-
-
- #Original formula
-
- 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
-
- _top = complex(_mass * width, 0.0)
-
- _bottom = complex((_mass**2 - q2), -_mass*gamma_j)
-
- com = _top/_bottom * scale
-
- r = abs(com)
-
- _phase = c.phase(com)
-
- _phase += phase
-
- x = c.cos(phase)*r
- y = c.sin(phase)*r
-
- com = complex(x,y)
-
- return self.total_scale_amp*com
-
-
- def add_resonance(self, _mass, width, phase, scale):
-
- #Adds the resonace to the pdf in form of a string (to be executed later)
-
- self.total_pdf_string += "+ self.resonance(q2,{0},{1},{2},{3})".format(_mass, width, phase, scale)
-
- def bifur_gauss(self, q2, mean, amp, sigma_L, sigma_R):
-
- q = np.sqrt(q2)
-
- if q < mean:
- sigma = sigma_L
- else:
- sigma = sigma_R
-
- _exp = np.exp(- (q-mean)**2 / (2 * sigma**2))
-
- dgamma = amp*_exp/(np.sqrt(2*np.pi))*2*(sigma_L*sigma_R)/(sigma_L+sigma_R)
-
- com = complex(dgamma, 0)
-
- return self.total_scale_amp*com
-
- def add_cusp(self, mean, amp, sigma_L, sigma_R):
-
- self.total_pdf_string += "+ self.bifur_gauss(q2,{0},{1},{2},{3})".format(mean, "cusp_amp", sigma_L, sigma_R)
-
- self.cusp_mean = mean
- self.cusp_sigma_L = sigma_L
- self.cusp_sigma_R = sigma_R
- self.cusp_amp = amp
-
-
- def normalize_pdf(self):
-
- self.total_scale_amp = 1.0
-
- x_scan = np.linspace(self.x_min, self.x_max, 10000)
-
- y_scan = []
-
- for i in x_scan:
- y_scan.append(self.total_pdf(i**2))
-
- _mean = np.mean(y_scan)
-
- self.total_scale_amp = 1.0/(self.x_max-self.x_min)*_mean
-
- self._mean = _mean * self.total_scale_amp
-
- def log_likelihood(self, x_part, cusp_amp):
-
- _sum = 0.0
-
- for i in x_part:
-
- _sum += np.log(self.total_pdf(i**2, cusp_amp = cusp_amp, scan = True))
-
- return _sum
-