def final_state_momenta(self, m2ab, m2bc):
"""
Calculate 4-momenta of final state tracks in a certain reference frame
(decay is in x-z plane, particle A moves along z axis)
m2ab, m2bc : invariant masses of AB and BC combinations
"""
m2ac = self.msqsum - m2ab - m2bc
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)
p4a = atfk.lorentz_vector(
atfk.vector(atfi.zeros(p_a), atfi.zeros(p_a), p_a),
atfi.sqrt(p_a**2 + self.ma2))
p4b = atfk.lorentz_vector(
atfk.vector(p_b * atfi.sqrt(1. - cos_theta_b**2), atfi.zeros(p_b),
-p_b * cos_theta_b), atfi.sqrt(p_b**2 + self.mb2))
p4c = atfk.lorentz_vector(
atfk.vector(-p_c * atfi.sqrt(1. - cos_theta_c**2), atfi.zeros(p_c),
-p_c * cos_theta_c), atfi.sqrt(p_c**2 + self.mc2))
return (p4a, p4b, p4c)
def rotate(v, angle, axis):
"""rotate vector around an arbitrary axis, from ROOT implementation
:param v:
:param angle:
:param axis:
"""
if (angle != atfi.zeros(angle)):
ll = norm(axis)
if (ll == atfi.zeros(ll)):
sys.exit('ERROR in rotate: rotation axis is zero')
else:
sa = atfi.sin(angle)
ca = atfi.cos(angle)
dx = x_component(axis) / ll
dy = y_component(axis) / ll
dz = z_component(axis) / ll
vx = x_component(v)
vy = y_component(v)
vz = z_component(v)
_vx = (ca+(1-ca)*dx*dx)*vx + ((1-ca)*dx*dy-sa*dz) * \
vy + ((1-ca)*dx*dz+sa*dy)*vz
_vy = ((1-ca)*dy*dx+sa*dz)*vx + (ca+(1-ca)*dy*dy) * \
vy + ((1-ca)*dy*dz-sa*dx)*vz
_vz = ((1-ca)*dz*dx-sa*dy)*vx + \
((1-ca)*dz*dy+sa*dx)*vy + (ca+(1-ca)*dz*dz)*vz
return vector(_vx, _vy, _vz)
else:
return v
def generate_rho(cuts, rnd):
"""
Generate random combinations of rho -> pi pi and a kaon track
"""
meanrhopt = cuts[5]
meankpt = cuts[0]
ptcut = cuts[2]
mrhogen = breit_wigner_random(rnd[:, 0], mrho, wrho)
ones = atfi.ones(rnd[:, 0])
p = atfk.two_body_momentum(mrhogen, mpi * ones, mpi * ones)
zeros = atfi.zeros(p)
mom = [
atfk.lorentz_vector(atfk.vector(p, zeros, zeros),
atfi.sqrt(p**2 + mpi**2)),
atfk.lorentz_vector(-atfk.vector(p, zeros, zeros),
atfi.sqrt(p**2 + mpi**2))
]
mom = generate_rotation_and_boost(mom, mrhogen, meanrhopt, ptcut, rnd[:,
2:8])
p4k = generate_4momenta(rnd[:, 8:11], meankpt, ptcut, atfi.const(mk))
p4pi1 = mom[0]
p4pi2 = mom[1]
return p4k, p4pi1, p4pi2
def generate_kstar(cuts, rnd):
"""
Generate random combinations of Kstar and a pion track
"""
meankstarpt = cuts[4]
meanpipt = cuts[1]
ptcut = cuts[2]
mkstargen = breit_wigner_random(rnd[:, 0], mkstar, wkstar)
ones = atfi.ones(rnd[:, 0])
p = atfk.two_body_momentum(mkstargen, mk * ones, mpi * ones)
zeros = atfi.zeros(p)
mom = [
atfk.lorentz_vector(atfk.vector(p, zeros, zeros),
atfi.sqrt(p**2 + mk**2)),
atfk.lorentz_vector(-atfk.vector(p, zeros, zeros),
atfi.sqrt(p**2 + mpi**2))
]
mom = generate_rotation_and_boost(mom, mkstargen, meankstarpt, ptcut,
rnd[:, 2:8])
p4k = mom[0]
p4pi1 = mom[1]
p4pi2 = generate_4momenta(rnd[:, 8:11], meanpipt, ptcut, atfi.const(mpi))
return p4k, p4pi1, p4pi2
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(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 axes_before_rotation(pb):
"""Calculate old (before rotation) axes in the frame aligned with the momentum vector pb
:param pb:
"""
z1 = unit_vector(
spatial_components(pb)) # New z-axis is in the direction of pb
eb = time_component(pb)
z0 = vector(atfi.zeros(eb), atfi.zeros(eb),
atfi.ones(eb)) # Old z-axis vector
x0 = vector(atfi.ones(eb), atfi.zeros(eb),
atfi.zeros(eb)) # Old x-axis vector
sp = scalar_product(z1, z0)
a0 = z0 - z1 * scalar(sp) # vector in z-pb plane perpendicular to z0
x1 = tf.where(tf.equal(sp, 1.0), x0, -unit_vector(a0))
y1 = vector_product(z1, x1) # New y-axis
x = vector(x_component(x1), x_component(y1), x_component(z1))
y = vector(y_component(x1), y_component(y1), y_component(z1))
z = vector(z_component(x1), z_component(y1), z_component(z1))
return (x, y, z)
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 four_momenta_from_helicity_angles(md, ma, mb, theta, phi):
"""Calculate the four-momenta of the decay products in D->AB in the rest frame of D
md: mass of D
ma: mass of A
mb: mass of B
theta: angle between A momentum in D rest frame and D momentum in its helicity frame
phi: angle of plane formed by A & B in D helicity frame
:param md:
:param ma:
:param mb:
:param theta:
:param phi:
"""
# Calculate magnitude of momentum in D rest frame
p = two_body_momentum(md, ma, mb)
# Calculate energy in D rest frame
Ea = atfi.sqrt(p**2 + ma**2)
Eb = atfi.sqrt(p**2 + mb**2)
# Construct four-momenta with A aligned with D in D helicity frame
Pa = lorentz_vector(vector(atfi.zeros(p), atfi.zeros(p), p), Ea)
Pb = lorentz_vector(vector(atfi.zeros(p), atfi.zeros(p), -p), Eb)
# rotate four-momenta
Pa = rotate_lorentz_vector(Pa, atfi.zeros(phi), -theta, -phi)
Pb = rotate_lorentz_vector(Pb, atfi.zeros(phi), -theta, -phi)
return Pa, Pb
def rotate(v, angle, axis):
"""
Rotate vector around an arbitrary axis. Uses ROOT implementation.
:param v: Input 3-vector
:param angle: Rotation angle
:param axis: 3-vector defining rotation axis
:returns: Rotated vector
"""
if angle != atfi.zeros(angle):
ll = norm(axis)
if ll == atfi.zeros(ll):
sys.exit("ERROR in rotate: rotation axis is zero")
else:
sa = atfi.sin(angle)
ca = atfi.cos(angle)
dx = x_component(axis) / ll
dy = y_component(axis) / ll
dz = z_component(axis) / ll
vx = x_component(v)
vy = y_component(v)
vz = z_component(v)
_vx = ((ca + (1 - ca) * dx * dx) * vx +
((1 - ca) * dx * dy - sa * dz) * vy +
((1 - ca) * dx * dz + sa * dy) * vz)
_vy = (((1 - ca) * dy * dx + sa * dz) * vx +
(ca + (1 - ca) * dy * dy) * vy +
((1 - ca) * dy * dz - sa * dx) * vz)
_vz = (((1 - ca) * dz * dx - sa * dy) * vx +
((1 - ca) * dz * dy + sa * dx) * vy +
(ca + (1 - ca) * dz * dz) * vz)
return vector(_vx, _vy, _vz)
else:
return v
def axes_after_rotation(pb, oldaxes=None):
"""Calculate new (rotated) axes aligned with the momentum vector pb
:param pb:
:param oldaxes: (Default value = None)
"""
z1 = unit_vector(
spatial_components(pb)) # New z-axis is in the direction of pb
eb = time_component(pb)
zeros = atfi.zeros(eb)
ones = atfi.ones(eb)
# Old z-axis vector
z0 = vector(zeros, zeros, ones) if oldaxes == None else oldaxes[2]
# Old x-axis vector
x0 = vector(ones, zeros, zeros) if oldaxes == None else oldaxes[0]
sp = scalar_product(z1, z0)
a0 = z0 - z1 * scalar(sp) # vector in z-pb plane perpendicular to z0
x1 = tf.where(scalar(tf.equal(sp, 1.)), x0, -unit_vector(a0))
y1 = vector_product(z1, x1) # New y-axis
return (x1, y1, z1)
def rotation_and_boost(ps, pb):
"""rotate and boost all momenta from the list ps to the rest frame of pb
After the rotation, the coordinate system is defined as:
z axis: direction of pb
y axis: perpendicular to the plane formed by the old z and pb
x axis: [y,z]
ps : list of Lorentz vectors to rotate and boost
pb : Lorentz vector defining the new frame
:param ps:
:param pb:
:returns: list of transformed Lorentz vectors
:rtype: ps1
"""
newaxes = axes_after_rotation(pb)
eb = time_component(pb)
zeros = atfi.zeros(eb)
# Boost vector in the rotated coordinates along z axis
boost = vector(zeros, zeros, -norm(spatial_components(pb)) / eb)
return nested_rotation_and_boost(ps, newaxes, boost)
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))
