Newer
Older
Master_thesis / scratch / 3583790 / raremodel.py
@saslie saslie on 2 Apr 2019 30 KB ...
  1. import ROOT
  2. #from ROOT import TTree, TFile, Double
  3. import numpy as np
  4. from pdg_const import pdg
  5. import matplotlib
  6. matplotlib.use("Qt5Agg")
  7. import matplotlib.pyplot as plt
  8. import pickle as pkl
  9. import sys
  10. import time
  11. from helperfunctions import display_time
  12. import cmath as c
  13.  
  14. mmu = pdg['muon_M']
  15. mb = pdg["bquark_M"]
  16. ms = pdg["squark_M"]
  17. mK = pdg["Ks_M"]
  18. mB = pdg["Bplus_M"]
  19.  
  20. class model:
  21. def __init__(self):
  22.  
  23. self.mmu = pdg['muon_M']
  24. self.mb = pdg["bquark_M"]
  25. self.ms = pdg["squark_M"]
  26. self.mK = pdg["Ks_M"]
  27. self.mB = pdg["Bplus_M"]
  28.  
  29. self.C7eff = pdg["C7eff"]
  30. self.C9eff = pdg["C9eff"]
  31. self.C10eff = pdg["C10eff"]
  32. #self.C1 = pdg["C1"]
  33. #self.C2 = pdg["C2"]
  34. #self.C3 = pdg["C3"]
  35. #self.C4 = pdg["C4"]
  36. self.GF = pdg["GF"] #Fermi coupling const.
  37. self.alpha_ew = pdg["alpha_ew"]
  38. self.Vts = pdg["Vts"]
  39. self.Vtb = pdg["Vtb"]
  40. self.x_min = 2*self.mmu
  41. self.x_max = (self.mB - self.mK) - 0.1
  42. self.total_pdf_string = "self.total_nonres(q2)"
  43. self.mode = ""
  44. self.total_scale_amp = 1.0
  45. self._mean = 0.0
  46. self.cusp_mean = 1
  47. self.cusp_sigma_L = 1
  48. self.cusp_sigma_R = 1
  49. self.cusp_amp = 0
  50.  
  51. def formfactor(self, q2, subscript):
  52. #check if subscript is viable
  53. if subscript != "0" and subscript != "+" and subscript != "T":
  54. raise ValueError('Wrong subscript entered, choose either 0, + or T')
  55. #get constants
  56. mh = self.mK
  57. mbstar0 = pdg["mbstar0"]
  58. mbstar = pdg["mbstar"]
  59. b0 = pdg["b0"]
  60. bplus = pdg["bplus"]
  61. bT = pdg["bT"]
  62. N = 3
  63. #some helperfunctions
  64. tpos = (self.mB - self.mK)**2
  65. tzero = (self.mB + self.mK)*(np.sqrt(self.mB)-np.sqrt(self.mK))**2
  66. z_oben = np.sqrt(tpos - q2) - np.sqrt(tpos - tzero)
  67. z_unten = np.sqrt(tpos - q2) + np.sqrt(tpos - tzero)
  68. z = z_oben/z_unten
  69. #calculate f0
  70. if subscript == "0":
  71. prefactor = 1/(1 - q2/(mbstar0**2))
  72. _sum = 0
  73. for i in range(N):
  74. _sum += b0[i]*(z**i)
  75. return prefactor * _sum
  76. #calculate f+ or fT
  77. else:
  78. prefactor = 1/(1 - q2/(mbstar**2))
  79. _sum = 0
  80. if subscript == "T":
  81. b = bT
  82. else:
  83. b = bplus
  84. for i in range(N):
  85. _sum += b[i] * (z**i - ((-1)**(i-N)) * (i/N) * z**N)
  86. return prefactor * _sum
  87. def axiv_nonres(self, q2):
  88. GF = self.GF
  89. alpha_ew = self.alpha_ew
  90. Vtb = self.Vtb
  91. Vts = self.Vts
  92. C10eff = self.C10eff
  93. mmu = self.mmu
  94. mb = self.mb
  95. ms = self.ms
  96. mK = self.mK
  97. mB = self.mB
  98. #Some helperfunctions
  99. beta = np.sqrt(abs(1. - 4. * self.mmu**2. / q2))
  100. 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)
  101. #prefactor in front of whole bracket
  102. prefactor1 = GF**2. *alpha_ew**2. * (abs(Vtb*Vts))**2 * kabs * beta / (128. * np.pi**5.)
  103. #left term in bracket
  104. bracket_left = 2./3. * kabs**2 * beta**2 * abs(C10eff*self.formfactor(q2, "+"))**2
  105. #middle term in bracket
  106. _top = 4. * mmu**2 * (mB**2 - mK**2) * (mB**2 - mK**2)
  107. _under = q2 * mB**2
  108. bracket_middle = _top/_under * abs(C10eff * self.formfactor(q2, "0"))**2
  109. return prefactor1 * (bracket_left + bracket_middle) * 2 * np.sqrt(q2)
  110.  
  111. def vec_nonres(self, q2):
  112. GF = self.GF
  113. alpha_ew = self.alpha_ew
  114. Vtb = self.Vtb
  115. Vts = self.Vts
  116. C7eff = self.C7eff
  117. C9eff = self.C9eff
  118. mmu = self.mmu
  119. mb = self.mb
  120. ms = self.ms
  121. mK = self.mK
  122. mB = self.mB
  123. #Some helperfunctions
  124. beta = np.sqrt(abs(1. - 4. * self.mmu**2. / q2))
  125. 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)
  126. #prefactor in front of whole bracket
  127. prefactor1 = GF**2. *alpha_ew**2. * (abs(Vtb*Vts))**2 * kabs * beta / (128. * np.pi**5.)
  128. #right term in bracket
  129. prefactor2 = kabs**2 * (1. - 1./3. * beta**2)
  130. abs_bracket = abs(C9eff * self.formfactor(q2, "+") + 2 * C7eff * (mb + ms)/(mB + mK) * self.formfactor(q2, "T"))**2
  131. bracket_right = prefactor2 * abs_bracket
  132. return prefactor1 * bracket_right * 2 * np.sqrt(q2)
  133. def total_nonres(self, q2):
  134. #Get constants
  135. GF = self.GF
  136. alpha_ew = self.alpha_ew
  137. Vtb = self.Vtb
  138. Vts = self.Vts
  139. C10eff = self.C10eff
  140. C9eff = self.C9eff
  141. C7eff = self.C7eff
  142. mmu = self.mmu
  143. mb = self.mb
  144. ms = self.ms
  145. mK = self.mK
  146. mB = self.mB
  147. #vector nonresonant part
  148. vec_nonres_part = self.vec_nonres(q2)
  149. #axial verctor nonresonant part including C7
  150. axiv_nonres_part = self.axiv_nonres(q2)
  151. #Complete term
  152. return self.total_scale_amp*complex(vec_nonres_part + axiv_nonres_part, 0.0)
  153.  
  154. def generate_points(self, set_size = 10000, x_min = 2* mmu, x_max = (mB - mK) - 0.1, save = True, verbose = 1, mode = "slim_points", resolution = 7.0, min_bin_scaling = 100):
  155. #Setup contants and variables
  156. if mode != "slim_points" and mode != "full_points" and mode != "fast_binned":
  157. raise ValueError('Wrong mode entered, choose either slim_points, full_points or fast_binned')
  158. self.mode = mode
  159. mB = self.mB
  160. mK = self.mK
  161. mmu = self.mmu
  162. #Range of function in MeV
  163. dq = np.linspace(x_min, x_max ,5000)
  164. x1 = 2500
  165. y1 = self.total_pdf(x1**2)
  166. x2 = 4000
  167. y2 = self.total_pdf(x2**2)
  168. #Translate to MeV**2
  169. dgamma_tot = []
  170.  
  171. for i in dq:
  172. dgamma_tot.append(self.total_pdf(i**2))
  173. dq2 = []
  174.  
  175. for i in dq:
  176. dq2.append(i**2)
  177. #Generate random points
  178. psi2s_mass, psi2s_width, psi2s_phase, psi2s_scale = pdg["psi2s"]
  179. _max = max(dgamma_tot)
  180. A1_x1 = (_max-y1)*x1
  181. A23_x1 = y1*x1
  182. A1_x2 = (_max-y2)*x2
  183. A23_x2 = y2*x2
  184. if mode == "slim_points":
  185. x_part = []
  186. y_part = []
  187.  
  188. print("Generating set of size {}...".format(int(set_size)))
  189.  
  190. #print(len(y))
  191.  
  192. #ROOT.TRandom1().SetSeed(0)
  193. if verbose != 0:
  194. verbose_calls = []
  195. j = 0
  196. o = 0
  197. while j < 100:
  198. j += verbose
  199. verbose_calls.append(int(set_size*j/100))
  200. start = time.time()
  201. counter = 0
  202. counter_x = 0
  203. while len(x_part) < set_size:
  204. counter += 1
  205. x = ROOT.TRandom1().Uniform(x_min, x_max)
  206. y = ROOT.TRandom1().Uniform(0, _max)
  207. if y < self.total_pdf(x**2):
  208. x_part.append(x)
  209. counter_x += 1
  210. #progress calls
  211. if verbose != 0:
  212. end = time.time()
  213. if o*verbose+verbose < 100 and counter%300 == 0:
  214. print(" {0}{1} completed".format(o*verbose+verbose, "%"))
  215. print(" Projected time left: {0}".format(display_time(int((end-start)*set_size/(len(x_part)+1)-(end-start)))))
  216. sys.stdout.write("\033[F")
  217. sys.stdout.write("\x1b[2K")
  218. sys.stdout.write("\033[F")
  219. sys.stdout.write("\x1b[2K")
  220. if o*verbose + verbose >=100:
  221. sys.stdout.write("\033[F")
  222. sys.stdout.write("\x1b[2K")
  223. print(" Time to generate set: {0}".format(display_time(int(end-start))))
  224. if len(x_part) == verbose_calls[o]:
  225. o += 1
  226. print(" {0} out of {1} particles chosen!".format(int(counter_x), counter))
  227. print(" Set generated!")
  228. #Save the set
  229. if save:
  230.  
  231. part_set = {"x_part" : x_part, "y_part": y_part, "counter_tot": counter, "counter_x": counter_x}
  232. pkl.dump( part_set, open("./data/set_{0}_range({1}-{2}).pkl".format(int(set_size),int(x_min), int(x_max)) , "wb" ) )
  233.  
  234. print(" Set saved!")
  235. print
  236. #returns all the chosen points (x_part, y_part) and all the random points generated (x, y)
  237. return x_part, y_part, counter
  238. if mode == "full_points":
  239. x = []
  240. y = []
  241.  
  242. x_part = []
  243. y_part = []
  244.  
  245. print("Generating set of size {}...".format(int(set_size)))
  246.  
  247. #print(len(y))
  248.  
  249. #ROOT.TRandom1().SetSeed(0)
  250. if verbose != 0:
  251. verbose_calls = []
  252. j = 0
  253. o = 0
  254. while j < 100:
  255. j += verbose
  256. verbose_calls.append(int(set_size*j/100))
  257. start = time.time()
  258. counter = 0
  259. counter_x = 0
  260. while len(x_part) < set_size:
  261. counter += 1
  262. x.append(ROOT.TRandom1().Uniform(x_min, x_max))
  263. y.append(ROOT.TRandom1().Uniform(0, _max))
  264. if y[-1] < self.total_pdf(x[-1]**2):
  265. x_part.append(x)
  266. y_part.append(y)
  267. counter_x += 1
  268. #progress calls
  269. if verbose != 0:
  270. end = time.time()
  271. if o*verbose+verbose < 100 and counter%300 == 0:
  272. print(" {0}{1} completed".format(o*verbose+verbose, "%"))
  273. print(" Projected time left: {0}".format(display_time(int((end-start)*set_size/(len(x_part)+1)-(end-start)))))
  274. sys.stdout.write("\033[F")
  275. sys.stdout.write("\x1b[2K")
  276. sys.stdout.write("\033[F")
  277. sys.stdout.write("\x1b[2K")
  278. if o*verbose + verbose >=100:
  279. sys.stdout.write("\033[F")
  280. sys.stdout.write("\x1b[2K")
  281. print(" Time to generate set: {0}".format(display_time(int(end-start))))
  282. if len(x_part) == verbose_calls[o]:
  283. o += 1
  284. print(" {0} out of {1} particles chosen!".format(len(x_part), counter))
  285. print(" Set generated!")
  286. #Save the set
  287. if save:
  288.  
  289. part_set = {"x_part" : x_part, "y_part": y_part, "counter": counter}
  290. raw_set = {"x" : x, "y": y, "counter": counter}
  291. pkl.dump( part_set, open("./data/set_{0}.pkl".format(int(set_size)) , "wb" ) )
  292. pkl.dump( part_set, open("./data/set_raw_{0}.pkl".format(int(set_size)) , "wb" ) )
  293.  
  294. print(" Sets saved!")
  295. print
  296. #returns all the chosen points (x_part, y_part) and all the random points generated (x, y)
  297. return x_part, y_part, counter
  298. if mode == "fast_binned":
  299.  
  300. nbins = int((x_max-x_min)/resolution)
  301. print("Generating set with {} bins...".format(nbins))
  302. dq = np.linspace(x_min, x_max ,nbins+1)
  303. bin_mean = []
  304. bin_true_height = []
  305. for i in range(len(dq)-1):
  306. a = dq[i]
  307. b = dq[i+1]
  308. c = (a+b)/2
  309. bin_mean.append(c)
  310. height = self.total_pdf(c**2)
  311. bin_true_height.append(height)
  312. _min = min(bin_true_height)
  313. for i in range(len(bin_true_height)):
  314. bin_true_height[i] = bin_true_height[i]/_min*min_bin_scaling
  315. start = time.time()
  316. bin_height = []
  317. for i in range(len(bin_mean)):
  318. bin_height.append(int(ROOT.TRandom1().Gaus(bin_true_height[i], np.sqrt(bin_true_height[i]))))
  319. #print(int(ROOT.TRandom1().Gaus(bin_true_height[i], np.sqrt(bin_true_height[i]))))
  320. #progress calls
  321. if verbose != 0:
  322. end = time.time()
  323. print(" Time to generate set: {0}s".format(end-start))
  324. print(" {0} bins simulated".format(nbins))
  325. print(" Set generated!")
  326. #Save the set
  327. if save:
  328. _ = 0
  329. for i in bin_height:
  330. _ += i
  331.  
  332. part_set = {"bin_mean" : bin_mean, "bin_height": bin_height, "nbins": nbins}
  333. pkl.dump( part_set, open("./data/binned_set_{0}bins_{1}part.pkl".format(nbins, _) , "wb" ) )
  334.  
  335. print(" Set saved!")
  336. print
  337. return bin_mean, bin_height, nbins
  338. def draw_plots(self, part_set, counter, mode,min_bin_scaling = 100, x_min = 2* mmu, x_max = (mB - mK) - 0.1, resolution = 7):
  339. if mode != "slim_points" and mode != "full_points" and mode != "fast_binned":
  340. raise ValueError('Wrong mode entered, choose either slim_points, full_points or fast_binned')
  341. if mode != self.mode:
  342. raise ValueError('self.mode and mode are different, set them to the same value')
  343. #Resolution based on experiment chosen to be ~7MeV
  344. #Setup contants and variables
  345. print("Generating plots")
  346. if mode == "fast_binned":
  347. mB = self.mB
  348. mK = self.mK
  349. mmu = self.mmu
  350. #Range of function in MeV
  351. dq = np.linspace(x_min, x_max ,5000)
  352. #Translate to MeV**2
  353. dq2 = []
  354.  
  355. for i in dq:
  356. dq2.append(i**2)
  357.  
  358. #calculate formfactors
  359. ff_plus = []
  360. ff_T = []
  361. ff_0 = []
  362.  
  363. for i in dq:
  364. ff_0.append(self.formfactor(i**2, "0"))
  365. ff_T.append(self.formfactor(i**2, "T"))
  366. ff_plus.append(self.formfactor(i**2, "+"))
  367. #calculate nonresonant
  368. dgamma_axiv_nonres = []
  369. dgamma_vec_nonres = []
  370. dgamma_tot = []
  371.  
  372. for i in dq:
  373. dgamma_axiv_nonres.append(self.axiv_nonres(i**2))
  374. dgamma_vec_nonres.append(self.vec_nonres(i**2))
  375. dgamma_tot.append(self.total_pdf(i**2))
  376. #Plot formfactors
  377. plt.plot(dq2, ff_0, label = "0")
  378. plt.plot(dq2, ff_T, label = "T")
  379. plt.plot(dq2, ff_plus, label = "+")
  380.  
  381. plt.grid()
  382. plt.title("Formfactors")
  383. plt.legend()
  384.  
  385. plt.savefig("./plots/fast_binned/ff.png")
  386.  
  387. print(" ff.png created")
  388. plt.clf()
  389. #Plot nonresonant part
  390. plt.plot(dq, dgamma_axiv_nonres, label = "axiv")
  391. plt.plot(dq, dgamma_vec_nonres, label = "vec")
  392.  
  393. plt.grid()
  394. plt.title("Nonresonant axial vector and vector parts")
  395. plt.legend()
  396.  
  397. plt.savefig("./plots/fast_binned/vec_axiv.png")
  398. print(" vec_axiv.png created")
  399.  
  400. plt.clf()
  401.  
  402. plt.plot(dq, dgamma_tot, label = "total")
  403.  
  404. plt.grid()
  405. plt.title("Total pdf")
  406.  
  407. plt.legend()
  408.  
  409. plt.savefig("./plots/fast_binned/tot.png")
  410. print(" tot.png created")
  411. #All pdfs
  412.  
  413. #print(test_x[1]**2 - self.x_min**2)
  414.  
  415. tot_y = []
  416. jpsi_y = []
  417. psi2s_y = []
  418. total_nonres_y = []
  419. cusp_y = []
  420.  
  421. jpsi_mass, jpsi_width, jpsi_phase, jpsi_scale = pdg["jpsi"]
  422.  
  423. psi2s_mass, psi2s_width, psi2s_phase, psi2s_scale = pdg["psi2s"]
  424.  
  425. for i in range(len(dq)):
  426. #print(i**2 - 4*(mmu**2))
  427. tot_y.append(abs(self.total_pdf(dq[i]**2)))
  428. jpsi_y.append(abs(self.resonance(dq[i]**2, jpsi_mass, jpsi_width, jpsi_phase, jpsi_scale)))
  429. psi2s_y.append(abs(self.resonance(dq[i]**2, psi2s_mass, psi2s_width, psi2s_phase, psi2s_scale)))
  430. total_nonres_y.append(abs(self.total_nonres(dq[i]**2)))
  431. cusp_y.append(abs(self.bifur_gauss(dq[i]**2, self.cusp_mean, self.cusp_amp, self.cusp_sigma_L, self.cusp_sigma_R )))
  432. #resonance(self, q2, _mass, width, phase, scale):
  433. #w[i] = np.sqrt(w[i])
  434. #print(test_y[i])
  435. plt.clf()
  436. plt.title("All pdfs")
  437. #plt.yscale("log")
  438.  
  439. plt.ylim(0, 1e-5)
  440.  
  441. plt.grid()
  442.  
  443. plt.plot(dq, tot_y, label = "total pdf")
  444. plt.plot(dq, jpsi_y, label = "jpsi")
  445. plt.plot(dq, psi2s_y, label = "psi2s")
  446. plt.plot(dq, total_nonres_y, label = "nonres")
  447. plt.plot(dq, cusp_y, label = "cusp")
  448.  
  449. plt.legend()
  450.  
  451. plt.savefig("./plots/fast_binned/pdf_and_parts.png")
  452. print(" pdf_and_parts.png created")
  453. #Create histo with pdf
  454.  
  455. #Translate to MeV**2
  456. plt.clf()
  457.  
  458. dq2 = []
  459.  
  460. for i in dq:
  461. dq2.append(i**2)
  462. dgamma_tot = []
  463.  
  464. for i in dq2:
  465. dgamma_tot.append(self.total_pdf(i))
  466.  
  467. _min = min(dgamma_tot)
  468.  
  469. for i in range(len(dgamma_tot)):
  470. dgamma_tot[i] = dgamma_tot[i]/_min*min_bin_scaling
  471. bin_mean, bin_height = part_set
  472.  
  473. nbins = counter
  474. plt.hist(bin_mean, bins=nbins, range=(self.x_min, self.x_max), weights = bin_height, label = "toy data binned")
  475.  
  476. plt.plot(dq, dgamma_tot, label = "pdf")
  477.  
  478. _sum = 0
  479.  
  480. for i in bin_height:
  481. _sum += i
  482.  
  483. #print(_max)
  484.  
  485. plt.grid()
  486.  
  487. plt.ylim(0, 2000)
  488.  
  489. plt.legend()
  490.  
  491. plt.title("{0} random points generated according to pdf ({1} particles)".format(len(bin_mean), _sum))
  492.  
  493. plt.savefig("./plots/fast_binned/histo.png")
  494.  
  495. print(" histo.png created")
  496. print(" All plots drawn \n")
  497. return
  498. else:
  499. mB = self.mB
  500. mK = self.mK
  501. mmu = self.mmu
  502. #Range of function in MeV
  503. dq = np.linspace(x_min, x_max ,5000)
  504. #Translate to MeV**2
  505. dq2 = []
  506.  
  507. for i in dq:
  508. dq2.append(i**2)
  509.  
  510. #calculate formfactors
  511. ff_plus = []
  512. ff_T = []
  513. ff_0 = []
  514.  
  515. for i in dq:
  516. ff_0.append(self.formfactor(i**2, "0"))
  517. ff_T.append(self.formfactor(i**2, "T"))
  518. ff_plus.append(self.formfactor(i**2, "+"))
  519. #calculate nonresonant
  520. dgamma_axiv_nonres = []
  521. dgamma_vec_nonres = []
  522. dgamma_tot = []
  523.  
  524. for i in dq:
  525. dgamma_axiv_nonres.append(self.axiv_nonres(i**2))
  526. dgamma_vec_nonres.append(self.vec_nonres(i**2))
  527. dgamma_tot.append(self.total_pdf(i**2))
  528. #Plot formfactors
  529. plt.plot(dq2, ff_0, label = "0")
  530. plt.plot(dq2, ff_T, label = "T")
  531. plt.plot(dq2, ff_plus, label = "+")
  532.  
  533. plt.grid()
  534. plt.title("Formfactors")
  535. plt.legend()
  536.  
  537. plt.savefig("./plots/points/ff.png")
  538.  
  539. print(" ff.png created")
  540. plt.clf()
  541. #Plot nonresonant part
  542. plt.plot(dq, dgamma_axiv_nonres, label = "axiv")
  543. plt.plot(dq, dgamma_vec_nonres, label = "vec")
  544.  
  545. plt.grid()
  546. plt.title("Nonresonant axial vector and vector parts")
  547. plt.legend()
  548.  
  549. plt.savefig("./plots/points/vec_axiv.png")
  550. print(" vec_axiv.png created")
  551.  
  552. plt.clf()
  553.  
  554. plt.plot(dq, dgamma_tot, label = "total")
  555.  
  556. plt.grid()
  557. plt.title("Total pdf")
  558.  
  559. plt.legend()
  560.  
  561. plt.savefig("./plots/points/tot.png")
  562. print(" tot.png created")
  563. #Particle set
  564. x_part, y_part = part_set
  565. set_size = len(x_part)
  566. #Plot generated generate_points
  567. #plt.clf()
  568.  
  569. #plt.plot(x_part, y_part, label = "Random points generated", marker = ".", linewidth = 0)
  570.  
  571. #plt.plot(dq, dgamma_tot, label = "pdf")
  572.  
  573. #plt.grid()
  574. #plt.title("Random points generated and pdf")
  575.  
  576. #plt.legend()
  577.  
  578. #plt.savefig("./plots/points/points_raw.png")
  579. #print(" points_raw.png created")
  580. #Histo unnormalized
  581. bins = int((x_max-x_min)/resolution)
  582. plt.clf()
  583. wheights = np.ones_like(x_part)
  584. _max = max(dgamma_tot)
  585. x1 = 2500
  586. y1 = self.total_pdf(x1**2)
  587. x2 = 4000
  588. y2 = self.total_pdf(x2**2)
  589. for i in range(len(wheights)):
  590. if x_part[i] < x1:
  591. wheights[i] = x1*y1/(x1*_max)
  592. elif x_part[i] > x2:
  593. wheights[i] = x2*y2/(x2*_max)
  594. _y, _x, _ = plt.hist(x_part, bins=bins, weights = wheights, range=(x_min, x_max), label = "toy data binned ({0} points)".format(sum(wheights)))
  595.  
  596. _mean_histo = float(np.mean(_y))
  597. plt.legend()
  598. plt.title("Binned toy data")
  599. plt.savefig("./plots/points/histo_raw.png")
  600. print(" histo_raw.png created")
  601. #Histo and pdf normailzed
  602. plt.clf()
  603.  
  604. for i in range(len(dgamma_tot)):
  605. dgamma_tot[i] = dgamma_tot[i]/(float(set_size)*_max * 2.0 * mmu / counter)
  606.  
  607. _mean = np.mean(dgamma_tot)
  608.  
  609. #Attempt for marked field of std-dev
  610. #dgamma_min = []
  611. #dgamma_plu = []
  612. #for i in range(len(dgamma_tot)):
  613. #dgamma_min.append(dgamma_tot[i]-np.sqrt(dgamma_tot[i]))
  614. #dgamma_plu.append(dgamma_tot[i]+np.sqrt(dgamma_tot[i]))
  615. #plt.plot(dq, dgamma_min, alpha = 0.5)
  616.  
  617. #plt.plot(dq, dgamma_plu, alpha = 0.5)
  618.  
  619. #plt.fill_between(dq, dgamma_min, dgamma_plu)
  620. #Plot histo
  621.  
  622. plt.hist(x_part, bins=bins, range=(x_min, x_max), weights = wheights/(_mean_histo/_mean), label = "toy data binned")
  623.  
  624. plt.plot(dq, dgamma_tot, label = "pdf")
  625.  
  626. #print(_max)
  627.  
  628. plt.grid()
  629.  
  630. plt.legend()
  631. plt.ylim(0, 1e-5)
  632. plt.title("{0} random points generated according to pdf".format(sum(wheights)))
  633.  
  634. plt.savefig("./plots/points/histo.png")
  635. print(" histo.png created")
  636. #All pdfs
  637.  
  638. #print(test_x[1]**2 - self.x_min**2)
  639.  
  640. tot_y = []
  641. jpsi_y = []
  642. psi2s_y = []
  643. total_nonres_y = []
  644. cusp_y = []
  645.  
  646. jpsi_mass, jpsi_width, jpsi_phase, jpsi_scale = pdg["jpsi"]
  647.  
  648. psi2s_mass, psi2s_width, psi2s_phase, psi2s_scale = pdg["psi2s"]
  649.  
  650. for i in range(len(dq)):
  651. #print(i**2 - 4*(mmu**2))
  652. tot_y.append(abs(self.total_pdf(dq[i]**2)))
  653. jpsi_y.append(abs(self.resonance(dq[i]**2, jpsi_mass, jpsi_width, jpsi_phase, jpsi_scale)))
  654. psi2s_y.append(abs(self.resonance(dq[i]**2, psi2s_mass, psi2s_width, psi2s_phase, psi2s_scale)))
  655. total_nonres_y.append(abs(self.total_nonres(dq[i]**2)))
  656. cusp_y.append(abs(self.bifur_gauss(dq[i]**2, self.cusp_mean, self.cusp_amp, self.cusp_sigma_L, self.cusp_sigma_R )))
  657. #resonance(self, q2, _mass, width, phase, scale):
  658. #w[i] = np.sqrt(w[i])
  659. #print(test_y[i])
  660. plt.clf()
  661. plt.title("All pdfs")
  662. #plt.yscale("log")
  663.  
  664. plt.ylim(0, 2*self._mean)
  665.  
  666. plt.grid()
  667.  
  668. plt.plot(dq, tot_y, label = "total pdf")
  669. plt.plot(dq, jpsi_y, label = "jpsi")
  670. plt.plot(dq, psi2s_y, label = "psi2s")
  671. plt.plot(dq, total_nonres_y, label = "nonres")
  672. plt.plot(dq, cusp_y, label = "cusp")
  673.  
  674. plt.legend()
  675.  
  676. plt.savefig("./plots/points/pdf_and_parts.png")
  677. print(" pdf_and_parts.png created")
  678. print(" All plots drawn \n")
  679. return
  680. def total_pdf(self, q2, cusp_amp = -1.0, scan = False):
  681. if scan:
  682. self.normalize_pdf()
  683. if cusp_amp == -1.0:
  684. cusp_amp = cusp_amp
  685. else:
  686. cusp_amp = self.cusp_amp
  687. #Calculate the pdf with the added resonances
  688. exec("_sum = abs({0})".format(self.total_pdf_string))
  689. return _sum
  690. def resonance(self, q2, _mass, width, phase, scale): #returns [real, imaginary]
  691. #calculate the resonance ---------------------------------------------> Formula correct?
  692. #if abs(np.sqrt(q2) - _mass) < 300:
  693. #return 0., 0.
  694. 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)
  695. #print(q2)
  696. #Teiler erweitert mit kompl. konj.
  697. #p = 0.5 * np.sqrt(q2 - 4*(mmu**2))
  698. #p0 = 0.5 * np.sqrt(_mass**2 - 4*mmu**2)
  699. #gamma_j = p / p0 * _mass /q2 * width
  700. #_top_im = - _mass**2 * width * gamma_j
  701. #_top_re = _mass * width * (_mass**2 - q2)
  702. #_bottom = (_mass**2 - q2)**2 + _mass**2 * gamma_j**2
  703. #real = _top_re/_bottom
  704. #imaginary = _top_im/_bottom
  705. #com = complex(real, imaginary) * scale
  706. #r = abs(com)
  707.  
  708. #_phase = c.phase(com)
  709.  
  710. #_phase += phase
  711.  
  712. #x = c.cos(phase)*r
  713. #y = c.sin(phase)*r
  714.  
  715. #com = complex(x,y)
  716. #Original formula
  717. p = 0.5 * np.sqrt(q2 - 4*(mmu**2))
  718. p0 = 0.5 * np.sqrt(_mass**2 - 4*mmu**2)
  719. gamma_j = p / p0 * _mass /q2 * width
  720. _top = complex(_mass * width, 0.0)
  721. _bottom = complex((_mass**2 - q2), -_mass*gamma_j)
  722. com = _top/_bottom * scale
  723. r = abs(com)
  724.  
  725. _phase = c.phase(com)
  726.  
  727. _phase += phase
  728.  
  729. x = c.cos(phase)*r
  730. y = c.sin(phase)*r
  731.  
  732. com = complex(x,y)
  733. return self.total_scale_amp*com
  734. def add_resonance(self, _mass, width, phase, scale):
  735. #Adds the resonace to the pdf in form of a string (to be executed later)
  736. self.total_pdf_string += "+ self.resonance(q2,{0},{1},{2},{3})".format(_mass, width, phase, scale)
  737. def bifur_gauss(self, q2, mean, amp, sigma_L, sigma_R):
  738. q = np.sqrt(q2)
  739. if q < mean:
  740. sigma = sigma_L
  741. else:
  742. sigma = sigma_R
  743. _exp = np.exp(- (q-mean)**2 / (2 * sigma**2))
  744. dgamma = amp*_exp/(np.sqrt(2*np.pi))*2*(sigma_L*sigma_R)/(sigma_L+sigma_R)
  745. com = complex(dgamma, 0)
  746. return self.total_scale_amp*com
  747. def add_cusp(self, mean, amp, sigma_L, sigma_R):
  748. self.total_pdf_string += "+ self.bifur_gauss(q2,{0},{1},{2},{3})".format(mean, "cusp_amp", sigma_L, sigma_R)
  749. self.cusp_mean = mean
  750. self.cusp_sigma_L = sigma_L
  751. self.cusp_sigma_R = sigma_R
  752. self.cusp_amp = amp
  753. def normalize_pdf(self):
  754. self.total_scale_amp = 1.0
  755. x_scan = np.linspace(self.x_min, self.x_max, 10000)
  756. y_scan = []
  757. for i in x_scan:
  758. y_scan.append(self.total_pdf(i**2))
  759. _mean = np.mean(y_scan)
  760. self.total_scale_amp = 1.0/(self.x_max-self.x_min)*_mean
  761. self._mean = _mean * self.total_scale_amp
  762. def log_likelihood(self, x_part, cusp_amp):
  763. _sum = 0.0
  764. for i in x_part:
  765. _sum += np.log(self.total_pdf(i**2, cusp_amp = cusp_amp, scan = True))
  766. return _sum