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))