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()
