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 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 get_gen_sample(sample='mu'):
den_columns = []
den_columns += ['Lambda_b0_TRUEP_E']
den_columns += ['Lambda_b0_TRUEP_X']
den_columns += ['Lambda_b0_TRUEP_Y']
den_columns += ['Lambda_b0_TRUEP_Z']
den_columns += ['Lambda_cplus_TRUEP_E']
den_columns += ['Lambda_cplus_TRUEP_X']
den_columns += ['Lambda_cplus_TRUEP_Y']
den_columns += ['Lambda_cplus_TRUEP_Z']
den_columns += [sample + 'minus_TRUEP_E']
den_columns += [sample + 'minus_TRUEP_X']
den_columns += [sample + 'minus_TRUEP_Y']
den_columns += [sample + 'minus_TRUEP_Z']
den_columns += ['nu_' + sample + '~_TRUEP_E']
den_columns += ['nu_' + sample + '~_TRUEP_X']
den_columns += ['nu_' + sample + '~_TRUEP_Y']
den_columns += ['nu_' + sample + '~_TRUEP_Z']
den_fname = '~/LbToLclnu_RunTwo/Selection/PID/FFs/GenMC/Lc' + sample.capitalize(
) + 'Nu_gen.root'
df_den = rpd.read_root(den_fname,
columns=den_columns,
key='MCDecayTreeTuple/MCDecayTree')
PLc_lab = atfk.lorentz_vector(
atfk.vector(df_den['Lambda_cplus_TRUEP_X'],
df_den['Lambda_cplus_TRUEP_Y'],
df_den['Lambda_cplus_TRUEP_Z']),
df_den['Lambda_cplus_TRUEP_E'])
Pl_lab = atfk.lorentz_vector(
atfk.vector(df_den[sample + 'minus_TRUEP_X'],
df_den[sample + 'minus_TRUEP_Y'],
df_den[sample + 'minus_TRUEP_Z']),
df_den[sample + 'minus_TRUEP_E'])
PNu_lab = atfk.lorentz_vector(
atfk.vector(df_den["nu_" + sample + "~_TRUEP_X"],
df_den["nu_" + sample + "~_TRUEP_Y"],
df_den["nu_" + sample + "~_TRUEP_Z"]),
df_den["nu_" + sample + "~_TRUEP_E"])
PLb_lab = PLc_lab + Pl_lab + PNu_lab
df_den['Lb_True_Q2'], df_den['Lb_True_Costhetal'] = get_phasespace_vars(
PLb_lab, PLc_lab, Pl_lab)
return df_den[['Lb_True_Q2', 'Lb_True_Costhetal']].to_numpy(),
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 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 generate_4momenta(rnd, meanpt, ptcut, m):
"""
Generate random 4-momenta according to specified mean Pt, minimum Pt, and mass of the particle, flat in eta and phi
"""
pt = generate_pt(rnd[:, 0], meanpt, ptcut, atfi.const(200.)) # Pt in GeV
eta = generate_eta(rnd[:, 1]) # Eta
phi = generate_phi(rnd[:, 2]) # Phi
theta = 2. * atfi.atan(atfi.exp(-eta))
p = pt / atfi.sin(theta) # Full momentum
e = atfi.sqrt(p**2 + m**2) # Energy
px = p * atfi.sin(theta) * atfi.sin(phi)
py = p * atfi.sin(theta) * atfi.cos(phi)
pz = p * atfi.cos(theta)
return atfk.lorentz_vector(atfk.vector(px, py, pz), e)
def kinematic_fit(mfit, moms):
"""
Kinematic fit to a fixed invariant mass for a multibody decay.
Returns the fitted mass and the list of final state 4-momenta.
"""
mcorr = atfk.mass(sum(moms))
for l in range(3):
dm2 = mcorr**2 - mfit**2
delta = momentum_scale(dm2, moms)
moms2 = []
for mom in moms:
m2 = atfk.mass(mom)**2
momvec = atfk.spatial_components(mom) * atfk.scalar(
1 + delta * momentum_resolution(atfk.p(mom)))
mom2 = atfk.lorentz_vector(momvec,
atfi.sqrt(m2 + atfk.norm(momvec)**2))
moms2 += [mom2]
moms = moms2
mcorr = atfk.mass(sum(moms))
return mcorr, moms
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))