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
Davide Lancierini