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
from zfit import ztf
#Create pdfs
def formfactor( 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
mK = ztf.constant(pdg['Ks_M'])
mbstar0 = ztf.constant(pdg["mbstar0"])
mbstar = ztf.constant(pdg["mbstar"])
b0 = ztf.constant(pdg["b0"])
bplus = ztf.constant(pdg["bplus"])
bT = ztf.constant(pdg["bT"])
mmu = ztf.constant(pdg['muon_M'])
mb = ztf.constant(pdg['bquark_M'])
ms = ztf.constant(pdg['squark_M'])
mB = ztf.constant(pdg['Bplus_M'])
#N comes from derivation in paper
N = 3
#some helperfunctions
tpos = (mB - mK)**2
tzero = (mB + mK)*(ztf.sqrt(mB)-ztf.sqrt(mK))**2
z_oben = ztf.sqrt(tpos - q2) - ztf.sqrt(tpos - tzero)
z_unten = ztf.sqrt(tpos - q2) + ztf.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
class resonance(zfit.pdf.ZPDF):
_N_OBS = 1 # dimension, can be omitted
_PARAMS = ['mass', 'width', 'phase'] # the name of the parameters
def _unnormalized_pdf(self, x):
x = x.unstack_x() # returns a list with the columns: do x, y, z = ztf.unstack_x(x) for 3D
q2 = tf.pow(x, 2)
_mass = self.params['mass']
width = self.params['width']
phase = self.params['phase']
mmu = ztf.constant(pdg['muon_M'])
p = 0.5 * ztf.sqrt(q2 - 4*(mmu**2))
p0 = 0.5 * ztf.sqrt(_mass**2 - 4*mmu**2)
gamma_j = tf.divide(p, q2) * _mass * width / p0
#Calculate the resonance
_top = tf.complex(_mass * width, ztf.constant(0.0))
_bottom = tf.complex(_mass**2 - q2, -_mass*gamma_j)
com = _top/_bottom
#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)
return com
class bifur_gauss(zfit.pdf.ZPDF):
_N_OBS = 1 # dimension, can be omitted
_PARAMS = ['mean', 'amp', 'sigma_L', 'sigma_R'] # the name of the parameters
def _unnormalized_pdf(self, x):
x = ztf.unstack_x(x) # returns a list with the columns: do x, y, z = ztf.unstack_x(x) for 3D
mean = self.params['mean']
amp = self.params['amp']
sigma_L = self.params['sigma_L']
sigma_R = self.params['sigma_R']
x_left = tf.where(x < mean, x, False)
x_right = tf.where(q >= mean, x, False)
#Calculate the exponential part of the cusp
_exp_left = ztf.exp(- tf.pow((x_left-mean),2) / (2 * sigma_L**2))
_exp_right = ztf.exp(- tf.pow((x_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/(ztf.sqrt(2*np.pi))*2*(sigma_L*sigma_R)/(sigma_L+sigma_R)
com = ztf.complex(dgamma, 0)
return com
def axiv_nonres(q2, scale_axiv):
GF = ztf.constant(pdg['GF'])
alpha_ew = ztf.constant(pdg['alpha_ew'])
Vtb = ztf.constant(pdg['Vtb'])
Vts = ztf.constant(pdg['Vts'])
C10eff = ztf.constant(pdg['C10eff'])
mmu = ztf.constant(pdg['muon_M'])
mb = ztf.constant(pdg['bquark_M'])
ms = ztf.constant(pdg['squark_M'])
mK = ztf.constant(pdg['Ks_M'])
mB = ztf.constant(pdg['Bplus_M'])
q2 = tf.pow(x, 2)
#Some helperfunctions
beta = ztf.sqrt(tf.abs(1. - 4. * mmu**2. / q2))
kabs = ztf.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*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 * formfactor(q2, "0")), 2)
#Note sqrt(q2) comes from derivation as we use q2 and plot q
return ztf.to_complex(scale_axiv * prefactor1 * (bracket_left + bracket_middle) * 2 *ztf.sqrt(q2))
def vec(q2, scale_vec):
scale = self.params['scale']
q2 = tf.pow(x, 2)
GF = ztf.constant(pdg['GF'])
alpha_ew = ztf.constant(pdg['alpha_ew'])
Vtb = ztf.constant(pdg['Vtb'])
Vts = ztf.constant(pdg['Vts'])
C7eff = ztf.constant(pdg['C7eff'])
C9eff = ztf.constant(pdg['C9eff'])
mmu = ztf.constant(pdg['muon_M'])
mb = ztf.constant(pdg['bquark_M'])
ms = ztf.constant(pdg['squark_M'])
mK = ztf.constant(pdg['Ks_M'])
mB = ztf.constant(pdg['Bplus_M'])
#Some helperfunctions
beta = ztf.sqrt(tf.abs(1. - 4. * mmu**2. / q2))
kabs = ztf.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 * formfactor(q2, "+") + 2 * C7eff * (mb + ms)/(mB + mK) * formfactor(q2, "T"))**2
bracket_right = prefactor2 * abs_bracket
#Note sqrt(q2) comes from derivation as we use q2 and plot q
return ztf.to_complex(scale_vec * prefactor1 * bracket_right * 2 * ztf.sqrt(q2))
class total_func(zfit.func.ZFunc):
_N_OBS = 1 # dimension, can be omitted
_PARAMS = ['jpsi_mass', 'jpsi_scale', 'jpsi_phase', 'jpsi_width',
'psi2s_mass', 'psi2s_scale', 'psi2s_phase', 'psi2s_width',
'scale_vec', 'scale_axiv',
'cusp_mass', 'sigma_L', 'sigma_R', 'cusp_scale'
] # the name of the parameters
def _func(self, x):
tot = tf.abs(vec_nonres + axiv_nonres)
return a
#Load data
x_min = 2*pdg['muon_M']
x_max = (pdg["Bplus_M"]-pdg["Ks_M"]-0.1)
obs = zfit.Space('q', limits = (x_min, x_max))
with open(r"./data/slim_points/slim_points_toy_0_range({0}-{1}).pkl".format(int(x_min), int(x_max)), "rb") as input_file:
part_set = pkl.load(input_file)
x_part = part_set['x_part']
x_part = x_part.astype('float64')
data = zfit.data.Data.from_numpy(array=x_part, obs=obs)
#Build pdf
#Old false nonres
nr_of_pdfs = 2
scale_test = zfit.Parameter("scale_test", ztf.constant(1.0))
frac1 = zfit.Parameter("frac1", 1/nr_of_pdfs)
frac2 = zfit.Parameter("frac2", 1/nr_of_pdfs)
axiv_nr = axiv_nonres(obs = obs)
vec_nr = vec_nonres(obs = obs, scale = scale_test)
#jpsi
jpsi_mass, jpsi_width, jpsi_phase, jpsi_scale = pdg["jpsi"]
jpsi_m = zfit.Parameter("jpsi_m", ztf.constant(jpsi_mass))
jpsi_w = zfit.Parameter("jpsi_w", ztf.constant(jpsi_width))
jpsi_p = zfit.Parameter("jpsi_p", ztf.constant(jpsi_phase))
# jpsi_res = resonance(obs = obs, mass = jpsi_m, width = jpsi_w, phase = jpsi_p)
#psi2s
psi2s_mass, psi2s_width, psi2s_phase, psi2s_scale = pdg["psi2s"]
total_pdf = total_func.as_pdf()
print(total_pdf.obs)
nll = zfit.loss.UnbinnedNLL(model=total_pdf, data=data)
minimizer = zfit.minimize.MinuitMinimizer()
result = minimizer.minimize(nll, )
param_errors = result.error()
for var, errors in param_errors.items():
print('{}: ^{{+{}}}_{{{}}}'.format(var.name, errors['upper'], errors['lower']))
print("Function minimum:", result.fmin)
#_____________________________________________________________________
#
# class c9pdf(BaseFunctor):
# __init__(c9pdf, param1, param2, param3, obs, name="c9pdf"):
# params =
#
# super()__init__(pdfs=c9pdf,...)
#
# def _unnormalized_pdf(..):
# c9pdf = self.pdfs[0]
# blabla... = bla c9pdf.pdf(x) ...
#
# #___________
#
# c9 = SumPDF(....)
# pdf = PDF(c9pdf=c9,...)
#
# #___________
#
# def _unnormalized_pdf(...):
# a, b, c = x.unstack_x()
# pdf1d = self.pdfs[0]
# pdf1d.pdf(x)
#
# #___________
#
# c9pdf = self.pdfs[0]
# c9eff = c9pdf.pdf(x)
