def final_state_momenta(self, x): """ Return final state momenta p(A1), p(A2), p(B1), p(B2) for the decay defined by the phase space vector x. The momenta are calculated in the D rest frame. """ ma1a2 = self.m_a1a2(x) mb1b2 = self.m_b1b2(x) ctha = self.cos_helicity_a(x) cthb = self.cos_helicity_b(x) phi = self.phi(x) p0 = atfk.two_body_momentum(self.md, ma1a2, mb1b2) pA = atfk.two_body_momentum(ma1a2, self.ma1, self.ma2) pB = atfk.two_body_momentum(mb1b2, self.mb1, self.mb2) zeros = atfi.zeros(pA) p3A = atfk.rotate_euler(Vector(zeros, zeros, pA), zeros, Acos(ctha), zeros) p3B = atfk.rotate_euler(Vector(zeros, zeros, pB), zeros, Acos(cthb), phi) ea = atfi.sqrt(p0 ** 2 + ma1a2 ** 2) eb = atfi.sqrt(p0 ** 2 + mb1b2 ** 2) v0a = atfk.vector(zeros, zeros, p0 / ea) v0b = atfk.vector(zeros, zeros, -p0 / eb) p4A1 = atfk.lorentz_boost(atfk.lorentz_vector(p3A, atfi.sqrt(self.ma1 ** 2 + pA ** 2)), v0a) p4A2 = atfk.lorentz_boost(atfk.lorentz_vector(-p3A, atfi.sqrt(self.ma2 ** 2 + pA ** 2)), v0a) p4B1 = atfk.lorentz_boost(atfk.lorentz_vector(p3B, atfi.sqrt(self.mb1 ** 2 + pB ** 2)), v0b) p4B2 = atfk.lorentz_boost(atfk.lorentz_vector(-p3B, atfi.sqrt(self.mb2 ** 2 + pB ** 2)), v0b) return (p4A1, p4A2, p4B1, p4B2)
def final_state_momenta(data): # Obtain the vectors of angles from the input tensor using the functions # provided by phasespace object cos_theta_jpsi = phsp.cos_theta1(data) cos_theta_phi = phsp.cos_theta2(data) phi = phsp.phi(data) # Rest-frame momentum of two-body Bs->Jpsi phi decay p0 = atfk.two_body_momentum(mb, mjpsi, mphi) # Rest-frame momentum of two-body Jpsi->mu mu decay pjpsi = atfk.two_body_momentum(mjpsi, mmu, mmu) # Rest-frame momentum of two-body phi->K K decay pphi = atfk.two_body_momentum(mphi, mk, mk) # Vectors of zeros and ones of the same size as the data sample # (needed to use constant values that do not depend on the event) zeros = atfi.zeros(phi) ones = atfi.ones(phi) # 3-vectors of Jpsi->mumu and phi->KK decays (in the corresponding rest frames), # rotated by the helicity angles p3jpsi = atfk.rotate_euler( atfk.vector(zeros, zeros, pjpsi * ones), zeros, atfi.acos(cos_theta_jpsi), zeros ) p3phi = atfk.rotate_euler( atfk.vector(zeros, zeros, pphi * ones), zeros, atfi.acos(cos_theta_phi), phi ) ejpsi = atfi.sqrt(p0 ** 2 + mjpsi ** 2) # Energy of Jpsi in Bs rest frame ephi = atfi.sqrt(p0 ** 2 + mphi ** 2) # Energy of phi in Bs rest frame v0jpsi = atfk.vector( zeros, zeros, p0 / ejpsi * ones ) # 3-vector of Jpsi in Bs rest frame v0phi = atfk.vector( zeros, zeros, -p0 / ephi * ones ) # 3-vector of phi in Bs rest frame # Boost momenta of final-state particles into Bs rest frame p4mu1 = atfk.lorentz_boost( atfk.lorentz_vector(p3jpsi, atfi.sqrt(mmu ** 2 + pjpsi ** 2) * ones), v0jpsi ) p4mu2 = atfk.lorentz_boost( atfk.lorentz_vector(-p3jpsi, atfi.sqrt(mmu ** 2 + pjpsi ** 2) * ones), v0jpsi ) p4k1 = atfk.lorentz_boost( atfk.lorentz_vector(p3phi, atfi.sqrt(mk ** 2 + pphi ** 2) * ones), v0phi ) p4k2 = atfk.lorentz_boost( atfk.lorentz_vector(-p3phi, atfi.sqrt(mk ** 2 + pphi ** 2) * ones), v0phi ) return (p4mu1, p4mu2, p4k1, p4k2)
def final_state_momenta(self, m2ab, m2bc, costhetaa, phia, phibc): """ Calculate 4-momenta of final state tracks in the 5D phase space m2ab, m2bc : invariant masses of AB and BC combinations (cos)thetaa, phia : direction angles of the particle A in the D reference frame phibc : angle of BC plane wrt. polarisation plane z x p_a """ thetaa = atfi.acos(costhetaa) m2ac = self.msqsum - m2ab - m2bc # Magnitude of the momenta p_a = atfk.two_body_momentum(self.md, self.ma, atfi.sqrt(m2bc)) p_b = atfk.two_body_momentum(self.md, self.mb, atfi.sqrt(m2ac)) p_c = atfk.two_body_momentum(self.md, self.mc, atfi.sqrt(m2ab)) cos_theta_b = (p_a * p_a + p_b * p_b - p_c * p_c) / (2. * p_a * p_b) cos_theta_c = (p_a * p_a + p_c * p_c - p_b * p_b) / (2. * p_a * p_c) # Fix momenta with p3a oriented in z (quantisation axis) direction p3a = atfk.vector(atfi.zeros(p_a), atfi.zeros(p_a), p_a) p3b = atfk.vector(p_b * Sqrt(1. - cos_theta_b**2), atfi.zeros(p_b), -p_b * cos_theta_b) p3c = atfk.vector(-p_c * Sqrt(1. - cos_theta_c**2), atfi.zeros(p_c), -p_c * cos_theta_c) # rotate vectors to have p3a with thetaa as polar helicity angle p3a = atfk.rotate_euler(p3a, atfi.const(0.), thetaa, atfi.const(0.)) p3b = atfk.rotate_euler(p3b, atfi.const(0.), thetaa, atfi.const(0.)) p3c = atfk.rotate_euler(p3c, atfi.const(0.), thetaa, atfi.const(0.)) # rotate vectors to have p3a with phia as azimuthal helicity angle p3a = atfk.rotate_euler(p3a, phia, atfi.const(0.), atfi.const(0.)) p3b = atfk.rotate_euler(p3b, phia, atfi.const(0.), atfi.const(0.)) p3c = atfk.rotate_euler(p3c, phia, atfi.const(0.), atfi.const(0.)) # rotate BC plane to have phibc as angle with the polarization plane p3b = atfk.rotate(p3b, phibc, p3a) p3c = atfk.rotate(p3c, phibc, p3a) # Define 4-vectors p4a = atfk.lorentz_vector(p3a, atfi.sqrt(p_a**2 + self.ma2)) p4b = atfk.lorentz_vector(p3b, atfi.sqrt(p_b**2 + self.mb2)) p4c = atfk.lorentz_vector(p3c, atfi.sqrt(p_c**2 + self.mc2)) return (p4a, p4b, p4c)
import sys import tensorflow as tf sys.path.append("../") import amplitf.interface as atfi import amplitf.kinematics as atfk atfi.set_seed(2) rndvec = tf.random.uniform([32, 3], dtype=atfi.fptype()) v = rndvec[:, 0] th = atfi.acos(rndvec[:, 1]) phi = (rndvec[:, 2] * 2 - 1) * atfi.pi() p = atfk.lorentz_vector( atfk.vector(atfi.zeros(v), atfi.zeros(v), atfi.zeros(v)), atfi.ones(v)) bp = atfk.lorentz_boost( p, atfk.rotate_euler(atfk.vector(v, atfi.zeros(v), atfi.zeros(v)), th, phi, atfi.zeros(v))) print(bp) print(atfk.mass(bp))