def gounaris_sakurai_lineshape(s, m, gamma, m_pi):
"""
Gounaris-Sakurai shape for rho->pipi
s : squared pipi inv. mass
m : rho mass
gamma : rho width
m_pi : pion mass
"""
m2 = m * m
m_pi2 = m_pi * m_pi
ss = atfi.sqrt(s)
ppi2 = (s - 4. * m_pi2) / 4.
p02 = (m2 - 4. * m_pi2) / 4.
p0 = atfi.sqrt(p02)
ppi = atfi.sqrt(ppi2)
hs = 2. * ppi / atfi.pi() / ss * atfi.log((ss + 2. * ppi) / 2. / m_pi)
hm = 2. * p0 / atfi.pi() / m * atfi.log((m + 2. * ppi) / 2. / m_pi)
dhdq = hm * (1. / 8. / p02 - 1. / 2. / m2) + 1. / 2. / atfi.pi() / m2
f = gamma * m2 / (p0**3) * (ppi2 * (hs - hm) - p02 * (s - m2) * dhdq)
gamma_s = gamma * m2 * (ppi**3) / s / (p0**3)
dr = m2 - s + f
di = ss * gamma_s
r = dr / (dr**2 + di**2)
i = di / (dr**2 + di**2)
return atfi.complex(r, i)
def unbinned_nll(pdf, integral):
"""
Return unbinned negative log likelihood graph for a PDF
pdf : PDF
integral : precalculated integral
"""
return -tf.reduce_sum(atfi.log(pdf / integral))
def eta(vector):
"""Return pseudorapidity component of the input Lorentz or 3-vector
:param vector: input vector (Lorentz or 3-vector)
:returns: vector of pseudorapidity components
"""
return -atfi.log(pt(vector) / 2. / z_component(vector))
def unbinned_weighted_nll(pdf, integral, weight_func):
"""
Return unbinned weighted negative log likelihood graph for a PDF
pdf : PDF
integral : precalculated integral
weight_func : weight function
"""
return -tf.reduce_sum(atfi.log(pdf / integral) * weight_func)
def generate_exp(rnd, x1, x2, alpha=None):
"""
Exponential random distribution with constant "alpha",
limited to the range x1 to x2
"""
if isinstance(x1, float): x1 = atfi.const(x1)
if isinstance(x2, float): x2 = atfi.const(x2)
if alpha is None or alpha == 0:
return uniform_random(rnd, x1, x2)
else:
if isinstance(alpha, float): alpha = atfi.const(alpha)
xi1 = atfi.exp(-x1 / alpha)
xi2 = atfi.exp(-x2 / alpha)
ximin = atfi.min(xi1, xi2)
ximax = atfi.max(xi1, xi2)
return atfi.abs(alpha * atfi.log(uniform_random(rnd, ximin, ximax)))
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
def unbinned_nll(pdf, integral):
return -tf.reduce_sum(atfi.log(pdf / integral))
def normal_random(rnd1, rnd2):
"""
Normal distribution from two random numbers
"""
return atfi.sqrt(-2. * atfi.log(rnd1)) * atfi.cos(2. * math.pi * rnd2)