def spin_rotation_angle(pa, pb, pc, bachelor=2):
"""Calculate the angle between two spin-quantisation axes for the 3-body D->ABC decay
aligned along the particle B and particle A.
pa, pb, pc : 4-momenta of the final-state particles
bachelor : index of the "bachelor" particle (0=A, 1=B, or 2=C)
:param pa:
:param pb:
:param pc:
:param bachelor: (Default value = 2)
"""
if bachelor == 2:
return atfi.const(0.)
pboost = lorentz_vector(
-spatial_components(pb) / scalar(time_component(pb)),
time_component(pb))
if bachelor == 0:
pa1 = spatial_components(lorentz_boost(pa, pboost))
pc1 = spatial_components(lorentz_boost(pc, pboost))
return atfi.acos(scalar_product(pa1, pc1) / norm(pa1) / norm(pc1))
if bachelor == 1:
pac = pa + pc
pac1 = spatial_components(lorentz_boost(pac, pboost))
pa1 = spatial_components(lorentz_boost(pa, pboost))
return atfi.acos(scalar_product(pac1, pa1) / norm(pac1) / norm(pa1))
return None
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(atfk.vector(zeros, zeros, pA), zeros, atfi.acos(ctha), zeros)
p3B = atfk.rotate_euler(atfk.vector(zeros, zeros, pB), zeros, atfi.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 helicity_angles_3body(pa, pb, pc):
"""Calculate 4 helicity angles for the 3-body D->ABC decay defined as:
theta_r, phi_r : polar and azimuthal angles of the AB resonance in the D rest frame
theta_a, phi_a : polar and azimuthal angles of the A in AB rest frame
:param pa:
:param pb:
:param pc:
"""
theta_r = atfi.acos(-z_component(pc) / norm(spatial_components(pc)))
phi_r = atfi.atan2(-y_component(pc), -x_component(pc))
pa_prime = lorentz_vector(
rotate_euler(spatial_components(pa), -phi_r,
atfi.pi() - theta_r, phi_r), time_component(pa))
pb_prime = lorentz_vector(
rotate_euler(spatial_components(pb), -phi_r,
atfi.pi() - theta_r, phi_r), time_component(pb))
w = time_component(pa) + time_component(pb)
pab = lorentz_vector(-(pa_prime + pb_prime) / scalar(w), w)
pa_prime2 = lorentz_boost(pa_prime, pab)
theta_a = atfi.acos(
z_component(pa_prime2) / norm(spatial_components(pa_prime2)))
phi_a = atfi.atan2(y_component(pa_prime2), x_component(pa_prime2))
return (theta_r, phi_r, theta_a, phi_a)
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 generate_rotation_and_boost(moms, minit, meanpt, ptcut, rnd):
"""
Generate 4-momenta of final state products according to 3-body phase space distribution
moms - initial particle momenta (in the rest frame)
meanpt - mean Pt of the initial particle
ptcut - miminum Pt of the initial particle
rnd - Auxiliary random tensor
"""
pt = generate_pt(rnd[:, 0], meanpt, ptcut, 200.) # Pt in GeV
eta = generate_eta(rnd[:, 1]) # Eta
phi = generate_phi(rnd[:, 2]) # Phi
theta = 2. * atfi.atan(atfi.exp(-eta)) # Theta angle
p = pt / atfi.sin(theta) # Full momentum
e = atfi.sqrt(p**2 + minit**2) # Energy
px = p * atfi.sin(theta) * atfi.sin(phi) # 3-momentum of initial particle
py = p * atfi.sin(theta) * atfi.cos(phi)
pz = p * atfi.cos(theta)
p4 = atfk.lorentz_vector(atfk.vector(px, py, pz),
e) # 4-momentum of initial particle
rotphi = uniform_random(rnd[:, 3], 0., 2 * atfi.pi())
rotpsi = uniform_random(rnd[:, 4], 0., 2 * atfi.pi())
rottheta = atfi.acos(uniform_random(rnd[:, 5], -1, 1.))
moms2 = []
for m in moms:
m1 = atfk.rotate_lorentz_vector(m, rotphi, rottheta, rotpsi)
moms2 += [atfk.boost_from_rest(m1, p4)]
return moms2
def m_prime_bc(self, sample):
"""
Square Dalitz plot variable m'
"""
mbc = atfi.sqrt(self.m2bc(sample))
return atfi.acos(2 * (mbc - math.sqrt(self.minbc)) /
(math.sqrt(self.maxbc) - math.sqrt(self.minbc)) -
1.) / math.pi
def m_prime_ab(self, sample):
"""
Square Dalitz plot variable m'
"""
mab = atfi.sqrt(self.m2ab(sample))
return atfi.acos(2 * (mab - math.sqrt(self.minab)) /
(math.sqrt(self.maxab) - math.sqrt(self.minab)) -
1.) / math.pi
def m_prime_ac(self, sample):
"""
Square Dalitz plot variable m'
"""
mac = atfi.sqrt(self.m2ac(sample))
return (
atfi.acos(2 * (mac - math.sqrt(self.minac)) /
(math.sqrt(self.maxac) - math.sqrt(self.minac)) - 1.0) /
math.pi)
def helicity_angles_4body(pa, pb, pc, pd):
"""Calculate 4 helicity angles for the 4-body E->ABCD decay defined as:
theta_ab, phi_ab : polar and azimuthal angles of the AB resonance in the E rest frame
theta_cd, phi_cd : polar and azimuthal angles of the CD resonance in the E rest frame
theta_ac, phi_ac : polar and azimuthal angles of the AC resonance in the E rest frame
theta_bd, phi_bd : polar and azimuthal angles of the BD resonance in the E rest frame
theta_ad, phi_ad : polar and azimuthal angles of the AD resonance in the E rest frame
theta_bc, phi_bc : polar and azimuthal angles of the BC resonance in the E rest frame
phi_ab_cd : azimuthal angle between AB and CD
phi_ac_bd : azimuthal angle between AC and BD
phi_ad_bc : azimuthal angle between AD and BC
:param pa:
:param pb:
:param pc:
:param pd:
"""
theta_r = atfi.acos(-z_component(pc) / norm(spatial_components(pc)))
phi_r = atfi.atan2(-y_component(pc), -x_component(pc))
pa_prime = lorentz_vector(
rotate_euler(spatial_components(pa), -phi_r,
atfi.pi() - theta_r, phi_r),
time_component(pa),
)
pb_prime = lorentz_vector(
rotate_euler(spatial_components(pb), -phi_r,
atfi.pi() - theta_r, phi_r),
time_component(pb),
)
w = time_component(pa) + time_component(pb)
pab = lorentz_vector(-(pa_prime + pb_prime) / scalar(w), w)
pa_prime2 = lorentz_boost(pa_prime, pab)
theta_a = atfi.acos(
z_component(pa_prime2) / norm(spatial_components(pa_prime2)))
phi_a = atfi.atan2(y_component(pa_prime2), x_component(pa_prime2))
return (theta_r, phi_r, theta_a, phi_a)
def spherical_angles(pb):
"""theta, phi : polar and azimuthal angles of the vector pb
:param pb:
"""
z1 = unit_vector(
spatial_components(pb)) # New z-axis is in the direction of pb
theta = atfi.acos(z_component(z1)) # Helicity angle
phi = atfi.atan2(y_component(pb), x_component(pb)) # phi angle
return (theta, phi)
def spherical_angles(pb):
"""
Return polar and azimuthal angles of the 3-vector or Lorentz vector
:param pb: Input vector
:returns: tuple (theta, phi), where theta is the polar
and phi the azimuthal angles
"""
z1 = unit_vector(
spatial_components(pb)) # Unit vector in the direction of pb
theta = atfi.acos(z_component(z1)) # Helicity angle
phi = atfi.atan2(y_component(pb), x_component(pb)) # phi angle
return (theta, phi)
def euler_angles(x1, y1, z1, x2, y2, z2):
"""Calculate Euler angles (phi, theta, psi in the ZYZ convention) which transform the coordinate basis (x1, y1, z1)
to the basis (x2, y2, z2). Both x1,y1,z1 and x2,y2,z2 are assumed to be orthonormal and right-handed.
:param x1:
:param y1:
:param z1:
:param x2:
:param y2:
:param z2:
"""
theta = atfi.acos(scalar_product(z1, z2))
phi = atfi.atan2(scalar_product(z1, y2), scalar_product(z1, x2))
psi = atfi.atan2(scalar_product(y1, z2), scalar_product(x1, z2))
return (phi, theta, psi)
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)
def random_rotation_and_boost(moms, rnd):
"""
Apply random boost and rotation to the list of 4-vectors
moms : list of 4-vectors
rnd : random array of shape (N, 6), where N is the length of 4-vector array
"""
pt = -5.0 * atfi.log(
rnd[:, 0]
) # Random pT, exponential distribution with mean 5 GeV
eta = rnd[:, 1] * 3.0 + 2.0 # Uniform distribution in pseudorapidity eta
phi = rnd[:, 2] * 2.0 * atfi.pi() # Uniform distribution in phi
theta = 2.0 * atfi.atan(atfi.exp(-eta)) # Theta angle is a function of eta
p = pt / atfi.sin(theta) # Full momentum
e = atfi.sqrt(p ** 2 + mb ** 2) # Energy of the Bs
px = (
p * atfi.sin(theta) * atfi.sin(phi)
) # 3-momentum of initial particle (Bs meson)
py = p * atfi.sin(theta) * atfi.cos(phi)
pz = p * atfi.cos(theta)
boost = atfk.lorentz_vector(
atfk.vector(px, py, pz), e
) # Boost vector of the Bs meson
rot_theta = atfi.acos(
rnd[:, 3] * 2.0 - 1.0
) # Random Euler rotation angles for Bs decay
rot_phi = rnd[:, 4] * 2.0 * atfi.pi()
rot_psi = rnd[:, 5] * 2.0 * atfi.pi()
# Apply rotation and boost to the momenta in input list
moms1 = []
for m in moms:
m1 = atfk.rotate_lorentz_vector(m, rot_phi, rot_theta, rot_psi)
moms1 += [atfk.boost_from_rest(m1, boost)]
return moms1
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))
def theta_prime_bc(self, sample):
"""
Square Dalitz plot variable theta'
"""
return atfi.acos(-self.cos_helicity_bc(sample)) / math.pi
