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, #phi_l, cos_theta_gamma, phi_gamma, #cos_theta_Kst, phi_Kst, 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 angles """ sin_theta_l=sqrt_(1-square_(cos_theta_l)) #phi_l=draw_random_phi(q2, mode) cos_phi_l = draw_random_ones_minusones(q2, mode)#one_ sin_phi_l = zero_#1.-cos_phi_l#zero_ sin_theta_gamma=sqrt_(1-square_(cos_theta_gamma)) cos_phi_gamma = cos_(phi_gamma) sin_phi_gamma = sin_(phi_gamma) cos_theta_Kst = draw_random_ones_minusones(q2, mode) sin_theta_Kst = zero_#1.-cos_theta_Kst #phi_Kst=draw_random_phi(q2, mode) cos_phi_Kst = one_#draw_random_ones_minusones(q2, mode) sin_phi_Kst = zero_#1.-cos_phi_Kst """ q-RF """ beta_l = dphi2_(sqrt_(q2), m_l, m_l) El = sqrt_(q2)/2 mod_l = El * beta_l p1_1 = mod_l*sin_theta_l*cos_phi_l p1_2 = mod_l*sin_theta_l*sin_phi_l p1_3 = mod_l*cos_theta_l p1_0 = El # 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 """ lambda_bar = lambda_f(sqrt_(mBbar2),mKst,sqrt_(q2)) beta_K = sqrt_(lambda_bar)/(mBbar2-q2+mKst**2) beta_q = -sqrt_(lambda_bar)/(mBbar2+q2-mKst**2) gamma_q = 1/sqrt_(1-square_(beta_q)) E_gamma = (mB**2-mBbar2-mgamma**2)/(2*sqrt_(mBbar2)) E_Kst = (mBbar2-q2+mKst**2)/(2*sqrt_(mBbar2)) E_q = (mBbar2+q2-mKst**2)/(2*sqrt_(mBbar2)) modk_Bbar = sqrt_(-mgamma**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_0, k_Bbar_1, k_Bbar_2, k_Bbar_3, ], axis =1) q_Bbar = conc_([ q_Bbar_0, q_Bbar_1, q_Bbar_2, q_Bbar_3, ], axis=1) pB_Bbar = conc_([ pB_Bbar_0, pB_Bbar_1, pB_Bbar_2, pB_Bbar_3, ], axis=1) # p1_Bbar_0 = El*gamma_q*(1 + beta_l*beta_q*cos_theta_Kst*cos_theta_l + beta_l*beta_q*cos_(phi_Kst-phi_l)*sin_theta_Kst*sin_theta_l) # p1_Bbar_1 = El*(beta_l*cos_phi_l*sin_theta_l + beta_l*(-1 + gamma_q)*cos_phi_Kst**2*cos_phi_l*sin_theta_Kst**2*sin_theta_l + cos_phi_Kst*sin_theta_Kst*(beta_q*gamma_q + beta_l*(-1 + gamma_q)*(cos_theta_Kst*cos_theta_l + sin_theta_Kst*sin_theta_l*sin_phi_Kst*sin_phi_l))) # p1_Bbar_2 = El*(sin_theta_Kst*(beta_q*gamma_q + beta_l*(-1 + gamma_q)*(cos_theta_Kst*cos_theta_l + cos_(phi_Kst-phi_l)*sin_theta_Kst*sin_theta_l))*sin_phi_Kst + beta_l*sin_theta_l*sin_phi_l) # p1_Bbar_3 = (El/2)*(2*beta_q*gamma_q*cos_theta_Kst + beta_l*(1 + gamma_q + (-1 + gamma_q)*cos_(2*theta_Kst))*cos_theta_l + beta_l*(-1 + gamma_q)*cos_(phi_Kst-phi_l)*sin_(2*theta_Kst)*sin_theta_l) # p2_Bbar_0 = -El*gamma_q*(-1 + beta_l*beta_q*cos_theta_Kst*cos_theta_l + beta_l*beta_q*cos_(phi_Kst-phi_l)*sin_theta_Kst*sin_theta_l) # p2_Bbar_1 = El*(-beta_l*cos_phi_l*sin_theta_l - beta_l*(-1 + gamma_q)*cos_phi_Kst**2*cos_phi_l*sin_theta_Kst**2*sin_theta_l + cos_phi_Kst*sin_theta_Kst*(beta_q*gamma_q - beta_l*(-1 + gamma_q)*(cos_theta_Kst*cos_theta_l + sin_theta_Kst*sin_theta_l*sin_phi_Kst*sin_phi_l))) # p2_Bbar_2 = El*(sin_theta_Kst*(beta_q*gamma_q - beta_l*(-1 + gamma_q)*(cos_theta_Kst*cos_theta_l + cos_(phi_Kst-phi_l)*sin_theta_Kst*sin_theta_l))*sin_phi_Kst - beta_l*sin_theta_l*sin_phi_l) # p2_Bbar_3 = (El/2)*(2*beta_q*gamma_q*cos_theta_Kst + 2*beta_l*(-1 - (-1 + gamma_q)*cos_theta_Kst**2)*cos_theta_l - beta_l*(-1 + gamma_q)*cos_(phi_Kst-phi_l)*sin_(2*theta_Kst)*sin_theta_l) # p1_Bbar = conc_([ # p1_Bbar_0, # p1_Bbar_1, # p1_Bbar_2, # p1_Bbar_3, # ],axis=1) # p2_Bbar = conc_([ # p2_Bbar_0, # p2_Bbar_1, # p2_Bbar_2, # p2_Bbar_3, # ],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_(p1, beta_Bbar) p2_Bbar = lorentz_boost_(p2, beta_Bbar) return pB_Bbar, pKst_Bbar, k_Bbar, q_Bbar, p1_Bbar, p2_Bbar, p1, p2