import numpy as np from math import pi from utils.kin_utils_edit import * import zfit ztf = zfit.ztf def momenta_map(q2, cos_theta_l, mB, mBbar2, mKst, m_l, mgamma, mode=None): if mode =='numpy': cos_ = np.cos sin_ = np.sin sqrt_ = np.sqrt square_ = np.square dphi2_ = dphi2_np zero_= np.zeros_like(q2) one_= np.ones_like(q2) pi_ = pi*np.ones_like(q2) conc_ = np.concatenate lorentz_boost_ = lorentz_boost_np n_events_ = q2.shape[0] lambda_f = lambda_function_np if mode =='tf': cos_ = tf.cos sin_ = tf.sin sqrt_ = ztf.sqrt square_ = ztf.square dphi2_ = dphi2_tf zero_ = tf.zeros_like(q2, dtype=tf.float64) one_ = tf.ones_like(q2, dtype=tf.float64) pi_ = pi*tf.ones_like(q2, dtype=tf.float64) conc_ = tf.stack lorentz_boost_ = lorentz_boost lambda_f = lambda_function """ Decay variables """ sin_theta_l = sqrt_(1-square_(cos_theta_l)) """ Integrated out variables """ # phi_gamma = draw_random_angle(q2,mode) # cos_phi_gamma = cos_(phi_gamma) # sin_phi_gamma = sin_(phi_gamma) # cos_theta_gamma = draw_random_costheta(q2, mode) # sin_theta_gamma = sqrt_(1-square_(cos_theta_gamma)) """ Mute variables """ phi_l = draw_random_angle(q2,mode) phi_Kst = draw_random_angle(q2,mode) cos_phi_l = cos_(phi_l) sin_phi_l = sin_(phi_l) cos_phi_Kst = cos_(phi_Kst) sin_phi_Kst = sin_(phi_Kst) cos_theta_Kst = draw_random_costheta(q2, mode) sin_theta_Kst = sqrt_(1-square_(cos_theta_Kst)) """ q-RF """ beta = dphi2_(sqrt_(q2), m_l, m_l) p1_1 = 0.5*sqrt_(q2)*beta*sin_theta_l*cos_phi_l p1_2 = 0.5*sqrt_(q2)*beta*sin_theta_l*sin_phi_l p1_3 = 0.5*sqrt_(q2)*beta*cos_theta_l p1_0 = 0.5*sqrt_(q2) p1 = conc_([ p1_0, p1_1, p1_2, p1_3], axis =1) p2 = conc_([ p1_0, -p1_1, -p1_2, -p1_3], axis =1) """ pBbar-RF """ mgamma_mass=mgamma lambda_bar = lambda_f(sqrt_(mBbar2),mKst,sqrt_(q2)) beta_K = sqrt_(lambda_bar)/(sqrt_(mBbar2)**2-q2+mKst**2) beta_q = -sqrt_(lambda_bar)/(sqrt_(mBbar2)**2+q2-mKst**2) # E_gamma = (mB**2-sqrt_(mBbar2)**2-mgamma**2)/(2*sqrt_(mBbar2)) E_Kst = (sqrt_(mBbar2)**2-q2+mKst**2)/(2*sqrt_(mBbar2)) E_q = (sqrt_(mBbar2)**2+q2-mKst**2)/(2*sqrt_(mBbar2)) # modk_Bbar = sqrt_(-mgamma_mass**2+E_gamma**2) pKst_Bbar_0 = E_Kst pKst_Bbar_1 = E_Kst*beta_K*(sin_theta_Kst*cos_phi_Kst) pKst_Bbar_2 = E_Kst*beta_K*(sin_theta_Kst*sin_phi_Kst) pKst_Bbar_3 = E_Kst*beta_K*cos_theta_Kst q_Bbar_0 = E_q q_Bbar_1 = E_q*beta_q*(sin_theta_Kst*cos_phi_Kst) q_Bbar_2 = E_q*beta_q*(sin_theta_Kst*sin_phi_Kst) q_Bbar_3 = E_q*beta_q*cos_theta_Kst # k_Bbar_0 = E_gamma # k_Bbar_1 = modk_Bbar*sin_theta_gamma*cos_phi_gamma # k_Bbar_2 = modk_Bbar*sin_theta_gamma*sin_phi_gamma # k_Bbar_3 = modk_Bbar*cos_theta_gamma # pB_Bbar_0 = sqrt_(mBbar2)+k_Bbar_0 # pB_Bbar_1 = k_Bbar_1 # pB_Bbar_2 = k_Bbar_2 # pB_Bbar_3 = k_Bbar_3 pKst_Bbar = conc_([ pKst_Bbar_0, pKst_Bbar_1, pKst_Bbar_2, pKst_Bbar_3], axis =1) # k_Bbar = conc_([k_Bbar_1, # k_Bbar_2, # k_Bbar_3, # k_Bbar_0], axis =1) q_Bbar = conc_([ q_Bbar_0, q_Bbar_1, q_Bbar_2, q_Bbar_3 ], axis=1) # pB_Bbar = conc_([ # pB_Bbar_1, # pB_Bbar_2, # pB_Bbar_3, # pB_Bbar_0 # ], axis=1) if mode =='numpy': beta_Bbar = q_Bbar/q_Bbar_0[:,0:1] if mode =='tf': beta_Bbar = q_Bbar/tf.expand_dims(q_Bbar_0,axis=-1) p1_Bbar = lorentz_boost_np(p1, beta_Bbar) p2_Bbar = lorentz_boost_np(p2, beta_Bbar) return pKst_Bbar, q_Bbar, p1_Bbar, p2_Bbar, p1, p2#pB_Bbar,