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
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)"
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(np.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. * (np.abs(Vtb*Vts))**2 * kabs * beta / (128. * np.pi**5.)
#left term in bracket
bracket_left = 2./3. * kabs**2 * beta**2 * np.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 * np.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(np.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. * (np.abs(Vtb*Vts))**2 * kabs * beta / (128. * np.pi**5.)
#right term in bracket
prefactor2 = kabs**2 * (1. - 1./3. * beta**2)
abs_bracket = np.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 vec_nonres_part + axiv_nonres_part
def generate_points(self, set_size = 10000, x_min = 2* mmu, x_max = (mB - mK) - 0.1, save = True, verbose = 1):
#Setup contants and variables
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
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)
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()
while len(x_part) < set_size:
x.append(ROOT.TRandom1().Uniform(x_min, x_max))
y.append(ROOT.TRandom1().Uniform(0, _max*1.05))
if y[-1] < self.total_pdf(x[-1]**2):
x_part.append(x[-1])
y_part.append(y[-1])
#progress calls
if verbose != 0:
end = time.time()
if o*verbose+verbose < 100 and len(x)%200 == 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), len(x)))
print("Set generated!")
#Save the set
if save:
part_set = {"x_part" : x_part, "y_part": y_part, "x": x, "y": y}
pkl.dump( part_set, open("set_{0}.pkl".format(int(set_size)) , "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, x, y
def draw_plots(self, part_set, x_min = 2* mmu, x_max = (mB - mK) - 0.1, resolution = 7):
#Resolution based on experiment chosen to be ~7MeV
#Setup contants and variables
print("Generating plots")
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("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("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("tot.png")
print("tot.png created")
#Particle set
x_part, y_part, x, y = part_set
set_size = len(x_part)
#Plot generated generate_points
plt.clf()
plt.plot(x, y, label = "total", marker = ".", linewidth = 0)
plt.plot(x_part, y_part, label = "chosen", marker = ".", linewidth = 0)
plt.plot(dq, dgamma_tot, label = "pdf")
plt.grid()
plt.title("Random points generated and pdf")
plt.legend()
plt.savefig("points_raw.png")
print("points_raw.png created")
#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.legend()
plt.title("Binned toy data")
plt.savefig("histo_raw.png")
print("histo_raw.png created")
_max = max(dgamma_tot)
#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 / float(len(x)))
_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 = np.ones_like(x_part)/(_mean_histo/_mean), label = "toy data binned")
plt.plot(dq, dgamma_tot, label = "pdf")
#print(_max)
plt.grid()
plt.legend()
plt.title("{0} random points generated according to pdf".format(len(x_part)))
plt.savefig("histo.png")
print("histo.png created")
print("All plots drawn \n")
return
def total_pdf(self, q2):
#Calculate the pdf with the added resonances
exec("_sum = " + self.total_pdf_string)
return _sum
def resonance(self, q2, _mass, width, phase, scale): #returns [real, imaginary]
#calculate the resonance ---------------------------------------------> Formula correct?
#if np.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)
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) + _mass**2 * gamma_j**2
real = _top_re/_bottom
imaginary = _top_im/_bottom
length = np.sqrt(real**2 + imaginary**2)
real = np.cos(phase)*length
imaginary = np.sin(phase)*length
return [scale * real, scale * imaginary]
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})[0]".format(_mass, width, phase, scale)
modl = model()
load_set = False
draw = True
set_size = 1e5
jpsi_mass, jpsi_width, jpsi_phase, jpsi_scale = pdg["jpsi"]
#modl.add_resonance(jpsi_mass, jpsi_width, jpsi_phase, jpsi_scale)
psi2s_mass, psi2s_width, psi2s_phase, psi2s_scale = pdg["psi2s"]
#modl.add_resonance(psi2s_mass, psi2s_width, psi2s_phase, psi2s_scale)
##modl.add_resonance(0.09, 3096.0, 1.02, -1.5)
#if load_set:
#with open(r"set.pkl", "rb") as input_file:
#set_dic = pkl.load(input_file)
#part_set = (set_dic["x_part"], set_dic["y_part"], set_dic["x"], set_dic["y"])
#else:
#part_set = modl.generate_points(set_size)
#if draw:
#modl.draw_plots(part_set = part_set)
#print(modl.total_pdf_string)
test_x = np.linspace(modl.x_min, modl.x_max, 1000)
#print(test_x[1]**2 - modl.x_min**2)
test_y = []
for i in range(len(test_x)):
#print(i**2 - 4*(mmu**2))
test_y.append(modl.resonance(test_x[i]**2, jpsi_mass, jpsi_width, jpsi_phase, jpsi_scale))
#resonance(self, q2, _mass, width, phase, scale):
#w[i] = np.sqrt(w[i])
#print(test_y[i])
plt.clf()
plt.plot(test_x, test_y)
plt.savefig("test.png")
print("Run finished")
