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)