import numpy as np from math import pi import tensorflow as tf import zfit ztf = zfit.ztf def metric_tensor_np(): """Metric tensor for Lorentz space (constant).""" return np.array([1.,-1.,-1.,-1.]) def metric_tensor(): """Metric tensor for Lorentz space (constant).""" return tf.constant([1., -1., -1., -1.], dtype=tf.float64) def scalar_product_3(v1, v2): """Calculate scalar product of two 3-vectors. Arguments: vec1: First vector. vec2: Second vector. """ return tf.reduce_sum(v1 * v2, axis=1) def scalar_product_4(v1, v2): """Calculate mass scalar for Lorentz 4-momentum. Arguments: vector: Input Lorentz momentum vector. """ return tf.reduce_sum(tf.multiply(tf.multiply(v1,v2),metric_tensor()), axis=-1, keepdims=True) def scalar_product_3_np(v1,v2): output = (v1*v2)[:,1:4].sum(axis=1, keepdims=True) return output def scalar_product_4_np(v1,v2): output = (v1*v2*metric_tensor_np()).sum(axis=1 ,keepdims=True) return output def get_costheta_l(p1, p2): num = scalar_product_3(p1, p2) den1= ztf.sqrt(scalar_product_3(p1, p1)) den2= ztf.sqrt(scalar_product_3(p2, p2)) costheta_l = num/(den1*den2) return costheta_l def get_costheta_l_np(p1, p2): num = scalar_product_3_np(p1, p2) den1= np.sqrt(scalar_product_3_np(p1, p1)) den2= np.sqrt(scalar_product_3_np(p2, p2)) costheta_l = num/(den1*den2) return costheta_l def beta_l(q2, m): beta_l = ztf.sqrt(1-4*(ztf.square(m)/q2)) return beta_l def beta_l_np(q2, m): beta_l = np.sqrt(1-4*(np.square(m)/q2)) return beta_l def lambda_function(m, c1, c2): #lambd = ztf.square(ztf.square(m)- ztf.square(c1)) -2*ztf.square(c2)*(ztf.square(m) + ztf.square(c1)) +ztf.square(ztf.square(c2)) lambd = (ztf.square(m) - ztf.square(c1 + c2))*(ztf.square(m) - ztf.square(c1 - c2)) return lambd def lambda_function_np(m, c1, c2): #lambd = np.square(np.square(m)- np.square(c1)) -2*np.square(c2)*(np.square(m) + np.square(c1)) +np.square(np.square(c2)) lambd = (np.square(m) - np.square(c1 + c2))*(np.square(m) - np.square(c1 - c2)) return lambd def dphi2_tf(m, c1, c2): lambd = lambda_function(m, c1, c2) out = ztf.sqrt(lambd)/(ztf.square(m)) return out def dphi2_np(m, c1, c2): lambd = lambda_function_np(m, c1, c2) out = np.sqrt(lambd)/(np.square(m)) return out def spatial_component(vector): return tf.gather(vector, indices=[1, 2, 3], axis=-1) def time_component(vector): return tf.gather(vector, indices=[0], axis=-1) def lorentz_vector(space, time): """Make a Lorentz vector from spatial and time components. Arguments: space: 3-vector of spatial components. time: Time component. """ return tf.concat([space, time], axis=-1) def lorentz_boost(fvector_to_boost, boost_vector): # boost_vector_time_component= time_component(boost_vector) #boost_vector_spatial = beta = spatial_component(boost_vector) beta2 = tf.expand_dims(scalar_product_3(beta,beta), axis=-1) gamma = 1./ztf.sqrt(1.-beta2) gamma_ = (gamma-1.)/beta2 ve = time_component(fvector_to_boost) vp = spatial_component(fvector_to_boost) ve2 = gamma*(ve + tf.expand_dims(scalar_product_3(vp,beta), axis=-1)) vp2 = vp + (gamma*ve + gamma_*tf.expand_dims(scalar_product_3(vp, beta), axis=-1))*beta boosted_fvector=lorentz_vector(ve2,vp2) return boosted_fvector def lorentz_boost_np(fvector_to_boost, boost_vector): beta = boost_vector beta2 = scalar_product_3_np(beta,beta) gamma = 1./np.sqrt(1.-beta2) gamma_ = (gamma-1.)/beta2 ve = fvector_to_boost[:,0].reshape(-1,1) vp = fvector_to_boost ve2 = gamma*(ve + scalar_product_3_np(fvector_to_boost,beta)) vp2 = vp + (gamma*ve + gamma_*scalar_product_3_np(fvector_to_boost, beta))*beta boosted_fvector = np.concatenate([ve2,vp2[:,1:4]],axis=1) return boosted_fvector def draw_random_angle(q2, mode): if mode=='numpy': random_phi_ = np.random.uniform(size=(q2.shape[0],1), low=-pi, high=pi, ) if mode=='tf': random_phi_ = tf.random.uniform(shape=tf.shape(q2), minval=-pi, maxval=pi, dtype=tf.float64) return random_phi_ def draw_random_costheta(q2, mode): if mode=='numpy': random_cosine_ = np.random.uniform(size=(q2.shape[0],1), low=-1.0, high=1.0, ) if mode=='tf': random_cosine_ = tf.random.uniform(shape=tf.shape(q2), minval=-1., maxval=1.0, dtype=tf.float64) return random_cosine_ def draw_random_ones_minusones(q2, mode): if mode=='numpy': n_events=q2.shape[0] mask = np.random.uniform(size=(n_events,1), low=0, high=1. )>0.5 random_one_=np.where(mask,1.,-1.) if mode=='tf': mask = tf.random.uniform(shape=tf.shape(q2), minval=0., maxval=1.0, dtype=tf.float64) random_one_ = tf.where(mask>0.5,tf.ones_like(q2),-tf.ones_like(q2)) return random_one_