import numpy as np
from pdg_const import pdg
import matplotlib
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 scipy.optimize import fminbound
from array import array as arr
import collections
from itertools import compress
import tensorflow as tf
import zfit as zf
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.list_of_BR = {}
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_val = collections.OrderedDict()
self.param_val_orig = collections.OrderedDict()
#helperconstants
self.x_min = 2*self.mmu
self.x_max = (self.mB - self.mK) - 0.1
self.total_pdf_list = []
self.total_pdf_names = []
self.mode = ""
self.param_val["total_scale_amp"] = 1.0
self.param_val["_mean"] = 0.0
self.reco_steps = 10000
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)*(tf.sqrt(self.mB)-tf.sqrt(self.mK))**2
z_oben = tf.sqrt(tpos - q2) - tf.sqrt(tpos - tzero)
z_unten = tf.sqrt(tpos - q2) + tf.sqrt(tpos - tzero)
z = tf.divide(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]*(tf.pow(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] * (tf.pow(z, i) - ((-1)**(i-N)) * (i/N) * tf.pow(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 = tf.sqrt(tf.abs(1. - 4. * self.mmu**2. / q2))
kabs = tf.sqrt(mB**2. + tf.pow(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. * (tf.abs(Vtb*Vts))**2. * kabs * beta / (128. * np.pi**5.)
#left term in bracket
bracket_left = 2./3. * kabs**2. * beta**2. * tf.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 * tf.pow(tf.abs(C10eff * self.formfactor(q2, "0")), 2)
#Note sqrt(q2) comes from derivation as we use q2 and plot q
return prefactor1 * (bracket_left + bracket_middle) * 2 * tf.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 = tf.sqrt(abs(1. - 4. * self.mmu**2. / q2))
kabs = tf.sqrt(mB**2. + tf.pow(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. * (tf.abs(Vtb*Vts))**2 * kabs * beta / (128. * np.pi**5.)
#right term in bracket
prefactor2 = kabs**2 * (1. - 1./3. * beta**2)
abs_bracket = tf.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 prefactor1 * bracket_right * 2 * tf.sqrt(q2)
def total_nonres(self, q2, parameters = 0, name = "", 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)
#Normalized factor
# for nr in self.nr_of_part_list:
# _sum += nr
#
# factor = self.nr_of_part_list[name]/_sum
factor = 1.0
#Complete term
if absolut:
return factor*(vec_nonres_part + axiv_nonres_part)
else:
return factor * tf.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
#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
set_sizes = np.random.poisson(set_size, nr_of_toys)
def total_pdf_neg(q2):
return -1*self.total_pdf(q2)
min_x = fminbound(total_pdf_neg, x_min**2, x_max**2)
maxi = self.total_pdf(np.array(min_x))
if np.abs(min_x-3096) < 10 or np.abs(min_x-3686) < 10:
maxi /= 1000
for toy in range(nr_of_toys):
set_size_intermed = 0
print("Generating toy {0} of size {1}...".format(toy+1, int(set_sizes[toy])))
x_part = []
_ = 0
while set_size_intermed < set_sizes[toy]:
x_part_raw = np.random.uniform(low = x_min, high = x_max, size = 2000000)
# print(x_part_raw[:40])
y_part_raw = np.random.uniform(low = 0.0, high = maxi, size = 2000000)
# print(y_part_raw[:40])
choose = np.where(y_part_raw < self.total_pdf(tf.pow(x_part_raw, 2)), True, False)
# print(choose[:40])
x_part_intermed = list(compress(x_part_raw, choose))
x_part += x_part_intermed
set_size_intermed += len(x_part_intermed)
x_part = x_part[:int(set_sizes[toy])]
x_part = np.array(x_part)
print(" Toy {0} of {1} generated!".format(toy+1, nr_of_toys))
#Save the set
if save:
part_set = {"x_part" : x_part}
pkl.dump( part_set, open("./data/slim_points/slim_points_toy_{0}_range({1}-{2}).pkl".format(toy, int(x_min), int(x_max)) , "wb" ) )
print(" Set saved!\n")
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":
set_sizes = np.random.poisson(nonres_set_size, nr_of_toys)
def total_pdf_neg(q2):
return -1*self.total_pdf(q2)
min_x = fminbound(total_pdf_neg, x_min**2, x_max**2)
maxi = self.total_pdf(np.array(min_x))
if np.abs(min_x-3096) < 10 or np.abs(min_x-3686) < 10:
maxi /= 1000
for toy in range(nr_of_toys):
set_size_intermed = 0
print("Generating true sized toy {0} of nonres size {1}...".format(toy+1, int(set_sizes[toy])))
x_part = []
_ = 0
while set_size_intermed < set_sizes[toy]:
x_part_raw = np.random.uniform(low = self.x_min, high = self.x_max, size = 2000000)
y_part_raw = np.random.uniform(low = 0.0, high = maxi, size = 2000000)
choose_pdf_tot = np.where(y_part_raw < self.total_pdf(tf.pow(x_part_raw, 2)), True, False)
choose_x_min = np.where(x_min < x_part_raw, True, False)
choose_x_max = np.where(x_max > x_part_raw, True, False)
choose = choose_pdf_tot * choose_x_max * choose_x_min
choose_nonres = np.where(y_part_raw < self.total_nonres(tf.pow(x_part_raw, 2)), True, False)
# print(choose[:40])
x_part_intermed = list(compress(x_part_raw, choose))
x_part += x_part_intermed
set_size_intermed += np.sum(choose_nonres)
x_part = x_part[:int(set_sizes[toy])]
x_part = np.array(x_part)
print(" Toy {0} of {1} generated!".format(toy+1, nr_of_toys))
#Save the set
if save:
part_set = {"x_part" : x_part}
pkl.dump( part_set, open("./data/slim_points/slim_points_toy_{0}_range({1}-{2}).pkl".format(toy, int(x_min), int(x_max)) , "wb" ) )
print(" Set saved!\n")
print
#returns all the chosen points (x_part, y_part) and all the random points generated (x, y)
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" and mode != "no_data":
raise ValueError('Wrong mode entered, choose either slim_points, fast_binned, true_data or no_data' )
if mode != self.mode:
raise ValueError('self.mode and mode are different, set them to the same value')
print("Generating plots")
if mode == "true_data":
folder = mode
elif mode == "fast_binned":
folder = mode
elif mode == "no_data":
folder = mode
else:
folder = "points"
#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 = tf.pow(dq, 2)
#calculate formfactors
ff_0 = self.formfactor(dq2, "0")
ff_T = self.formfactor(dq2, "T")
ff_plus = self.formfactor(dq2, "+")
#calculate nonresonant
dgamma_axiv_nonres = self.axiv_nonres(dq2)
dgamma_vec_nonres = self.vec_nonres(dq2)
dgamma_total_nonres = self.total_nonres(dq2, absolut = True)
dgamma_tot = self.total_pdf(dq2)
#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/{0}/ff.png".format(folder))
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/{0}/vec_axiv.png".format(folder))
print(" vec_axiv.png created")
#Plot all pdfs and parts
jpsi_mass, jpsi_width, jpsi_phase, jpsi_scale = pdg["jpsi"]
psi2s_mass, psi2s_width, psi2s_phase, psi2s_scale = pdg["psi2s"]
tot_y = self.total_pdf(dq2)
plt.clf()
_ = 0
for pdf in self.total_pdf_list:
name = self.total_pdf_names[_]
_pdf = pdf(dq2, param_val = self.param_val, absolut = True)
plt.plot(dq, _pdf, label = name)
_ += 1
if x_min < 3096 < x_max:
plt.ylim(0, 4e-2)
plt.plot(dq, tot_y, label = "total pdf")
prepare_plot("All pdfs")
# plt.ylim(0, 2*self.param_val[1])
plt.savefig("./plots/{0}/pdf_and_parts.png".format(folder))
if mode == "no_data":
print(" pdf_and_parts.png created\n")
else:
print(" pdf_and_parts.png created")
#Create histo with pdf
if mode == "fast_binned":
#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[
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
elif mode == "slim_points" or mode == "true_data":
#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)
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")
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
_sum = 0
for pdf in self.total_pdf_list:
_sum += pdf(q2, param_val = self.param_val)
return tf.abs(_sum)
def resonance(self, q2, parameters, name = "", absolut = False): #returns complex value
_mass, width, phase, scale = parameters[name + " mass"], parameters[name + " width"], parameters[name + " phase"], parameters[name + " scale"]
#Helper variables
p = 0.5 * tf.sqrt(q2 - 4*(mmu**2))
p0 = 0.5 * tf.sqrt(_mass**2 - 4*mmu**2)
gamma_j = tf.divide(p, q2) * _mass * width / p0
#Calculate the resonance
_top = tf.complex(_mass * width, 0.0)
_bottom = tf.complex(_mass**2 - q2, -_mass*gamma_j)
com = _top/_bottom * scale
#Rotate by the phase
r = tf.abs(com)
_phase = tf.angle(com)
_phase += phase
x = tf.cos(phase)*r
y = tf.sin(phase)*r
com = tf.complex(x, y)
if absolut:
return tf.abs(com)
else:
return com
def add_resonance(self, _mass, width, phase, scale, name, BR): #adds string to total_pdf and adds constant places for fit
#Adds the resonace to the pdf in form of a string (to be executed later)
self.param_val[name + " mass"] = _mass
self.param_val[name + " width"] = width
self.param_val[name + " phase"] = phase
self.param_val[name + " scale"] = scale
self.list_of_BR[name] = BR
self.param_val_orig = self.param_val
def reso_func(q2, param_val, absolut = False):
return self.resonance(q2, parameters = param_val, name = name, absolut = absolut)
self.total_pdf_names += [name]
self.total_pdf_list += [reso_func]
return
def bifur_gauss(self, q2, parameters, name = "", absolut = False): #returns complex value
mean, scale, sigma_L, sigma_R = parameters[name + " mass"], parameters[name + " scale"], parameters[name + " sigma_L"], parameters[name + " sigma_R"]
#Get q out of q2
q = tf.sqrt(q2)
#Choose the right sigma depending on side of the cusp
q_left = tf.where(q < mean, q, False)
q_right = tf.where(q >= mean, q, False)
#Calculate the exponential part of the cusp
_exp_left = tf.exp(- tf.pow((q_left-mean),2) / (2 * sigma_L**2))
_exp_right = tf.exp(- tf.pow((q_right-mean),2) / (2 * sigma_R**2))
#Scale so the total area under curve is 1 and the top of the cusp is continuous
_exp = _exp_left + _exp_right
dgamma = scale*_exp/(tf.sqrt(2*np.pi))*2*(sigma_L*sigma_R)/(sigma_L+sigma_R)
com = tf.complex(dgamma, 0)
# _sum = 0
#
# factor = 1.0
# for nr in self.nr_of_part_list:
# _sum += nr
#
# factor = self.nr_of_part_list[name]/_sum
# com *= factor
if absolut:
return tf.abs(com)
else:
return com
def add_cusp(self, mean, amp, sigma_L, sigma_R, BR): #adds string to total_pdf and adds constant places for fit
name = "cusp"
#Add cusp to total_pdf
def reso_func(q2, param_val, absolut = False):
return self.bifur_gauss(q2, parameters = param_val, name = name, absolut = absolut)
self.total_pdf_list += [reso_func]
self.total_pdf_names += [name]
#Save the variables of the cusp -> Used for fitting
self.param_val[name + " mass"] = mean
self.param_val[name + " sigma_L"] = sigma_L
self.param_val[name + " sigma_R"] = sigma_R
self.param_val[name + " scale"] = amp
self.list_of_BR[name] = BR
self.param_val_orig = self.param_val
def integrate_pdf(self, pdf_name, x_min = None, x_max = None, limit = 3000): #Normalizes pdf with a global factor in front, saves mean and global factor
if not x_min:
x_min = self.x_min
if not x_max:
x_max = self.x_max
if pdf_name == "total":
area, err = integrate.quad(self.total_pdf, x_min**2, x_max**2, limit = limit)
else:
_ = 0
for pdf in self.total_pdf_list:
name = self.total_pdf_names[_]
if name == pdf_name:
area, err = integrate.quad(lambda q2: pdf(q2, param_val = self.param_val, absolut = True), x_min**2, x_max**2, limit = limit)
_ += 1
if err/area>0.05:
print("Big error on integration of {0}, increasing the limit may help!")
return area
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)")
dic = self.param_val_orig
i = 0
else:
print("List of parameters (fitted values)")
dic = self.param_val
i = 0
print(" Nr. Name Value")
for key in dic.keys():
i += 1
print(" {0}. {1}: {2}".format(i, key, dic[key]))
print("")
def add_nonres(self, BR):
name = "nonres"
def nonres_func(q2, param_val, absolut = False):
return self.total_nonres(q2, parameters = param_val, name = name, absolut = absolut)
self.total_pdf_list += [nonres_func]
self.list_of_BR[name] = BR
self.total_pdf_names += [name]
def neg_log_likelihood(self, q2):
# self.normalize_pdf(verbose = 0)
ll = -1*tf.reduce_sum(tf.log(self.total_pdf(q2)))
return ll
def C9eff(self, q2):
C9eff = pdg["C9eff"]
for pdf in self.total_pdf_list:
C9eff += pdf(q2, param_val = self.param_val)
return C9eff
