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 inside(self, x):
"""
Check if the point x=(m2ab, m2bc) is inside the phase space
"""
m2ab = self.m2ab(x)
m2bc = self.m2bc(x)
mab = atfi.sqrt(m2ab)
inside = tf.logical_and(tf.logical_and(tf.greater(m2ab, self.minab), tf.less(m2ab, self.maxab)),
tf.logical_and(tf.greater(m2bc, self.minbc), tf.less(m2bc, self.maxbc)))
if self.macrange:
m2ac = self.msqsum - m2ab - m2bc
inside = tf.logical_and(inside, tf.logical_and(tf.greater(m2ac, self.macrange[0]**2), tf.less(m2ac, self.macrange[1]**2)))
if self.symmetric:
inside = tf.logical_and(inside, tf.greater(m2bc, m2ab))
eb = (m2ab - self.ma2 + self.mb2)/2./mab
ec = (self.md2 - m2ab - self.mc2)/2./mab
p2b = eb**2 - self.mb2
p2c = ec**2 - self.mc2
inside = tf.logical_and(inside, tf.logical_and(tf.greater(p2c, 0), tf.greater(p2b, 0)))
pb = atfi.sqrt(p2b)
pc = atfi.sqrt(p2c)
e2bc = (eb+ec)**2
m2bc_max = e2bc - (pb - pc)**2
m2bc_min = e2bc - (pb + pc)**2
return tf.logical_and(inside, tf.logical_and(tf.greater(m2bc, m2bc_min), tf.less(m2bc, m2bc_max)))
def _with_form_factor(self, dataset: dict) -> tf.Tensor:
inv_mass_squared = dataset[self._dynamics_props.inv_mass_name]
inv_mass = atfi.sqrt(inv_mass_squared)
mass0 = self._dynamics_props.resonance_mass
gamma0 = self._dynamics_props.resonance_width
m_a = atfi.sqrt(dataset[self._dynamics_props.inv_mass_name_prod1])
m_b = atfi.sqrt(dataset[self._dynamics_props.inv_mass_name_prod2])
meson_radius = self._dynamics_props.meson_radius
l_orbit = self._dynamics_props.orbit_angular_momentum
q_squared = two_body_momentum_squared(inv_mass, m_a, m_b)
q0_squared = two_body_momentum_squared(mass0, m_a, m_b)
ff2 = blatt_weisskopf_ff_squared(q_squared, meson_radius, l_orbit)
ff02 = blatt_weisskopf_ff_squared(q0_squared, meson_radius, l_orbit)
width = gamma0 * (mass0 / inv_mass) * (ff2 / ff02)
# So far its all in float64,
# but for the sqrt operation it has to be converted to complex
width = atfi.complex(width, tf.constant(
0.0, dtype=tf.float64)) * atfi.sqrt(
atfi.complex(
(q_squared / q0_squared),
tf.constant(0.0, dtype=tf.float64),
))
return relativistic_breit_wigner(
inv_mass_squared, mass0, width) * atfi.complex(
mass0 * gamma0 * atfi.sqrt(ff2), atfi.const(0.0))
def from_square_dalitz_plot(self, mprimeac, thprimeac):
"""
sample: Given mprimeac and thprimeac, returns 2D tensor for (m2ab, m2bc).
Make sure you don't pass in sqDP corner points as they lie outside phsp.
"""
m2AC = (
0.25 *
(self.maxac**0.5 * atfi.cos(math.pi * mprimeac) + self.maxac**0.5 -
self.minac**0.5 * atfi.cos(math.pi * mprimeac) + self.minac**0.5)
**2)
m2AB = (
0.5 *
(-(m2AC**2) + m2AC * self.ma**2 + m2AC * self.mb**2 +
m2AC * self.mc**2 + m2AC * self.md**2 - m2AC * atfi.sqrt(
(m2AC *
(m2AC - 2.0 * self.ma**2 - 2.0 * self.mc**2) + self.ma**4 -
2.0 * self.ma**2 * self.mc**2 + self.mc**4) / m2AC) *
atfi.sqrt(
(m2AC *
(m2AC - 2.0 * self.mb**2 - 2.0 * self.md**2) + self.mb**4 -
2.0 * self.mb**2 * self.md**2 + self.md**4) / m2AC) *
atfi.cos(math.pi * thprimeac) - self.ma**2 * self.mb**2 +
self.ma**2 * self.md**2 + self.mb**2 * self.mc**2 -
self.mc**2 * self.md**2) / m2AC)
m2BC = self.msqsum - m2AC - m2AB
return tf.stack([m2AB, m2BC], axis=1)
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 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 cos_helicity_angle_dalitz(m2ab, m2bc, md, ma, mb, mc):
"""Calculate cosine of helicity angle for a set of two Dalitz plot variables
:param m2ab: 1st Dalitz plot variable (squared invariant mass of the AB combination)
:param m2bc: 2nd Dalitz plot variable (squared invariant mass of the BC combination)
:param md: Mass of the decaying particle D
:param ma: Mass of the decay product A
:param mb: Mass of the decay product B
:param mc: Mass of the decay product C
:returns: Cosine of the helicity angle of the AB combination
(of the angle between AB direction in D rest frame
and A direction in AB rest frame).
"""
md2 = md**2
ma2 = ma**2
mb2 = mb**2
mc2 = mc**2
m2ac = md2 + ma2 + mb2 + mc2 - m2ab - m2bc
mab = atfi.sqrt(m2ab)
mac = atfi.sqrt(m2ac)
mbc = atfi.sqrt(m2bc)
p2a = 0.25 / md2 * (md2 - (mbc + ma)**2) * (md2 - (mbc - ma)**2)
p2b = 0.25 / md2 * (md2 - (mac + mb)**2) * (md2 - (mac - mb)**2)
p2c = 0.25 / md2 * (md2 - (mab + mc)**2) * (md2 - (mab - mc)**2)
eb = (m2ab - ma2 + mb2) / 2.0 / mab
ec = (md2 - m2ab - mc2) / 2.0 / mab
pb = atfi.sqrt(eb**2 - mb2)
pc = atfi.sqrt(ec**2 - mc2)
e2sum = (eb + ec)**2
m2bc_max = e2sum - (pb - pc)**2
m2bc_min = e2sum - (pb + pc)**2
return (m2bc_max + m2bc_min - 2.0 * m2bc) / (m2bc_max - m2bc_min)
def model(x):
# Get phase space variables
cosThetaK = phsp.coordinate(x, 0)
cosThetaL = phsp.coordinate(x, 1)
phi = phsp.coordinate(x, 2)
# Derived quantities
sinThetaK = atfi.sqrt(1.0 - cosThetaK * cosThetaK)
sinThetaL = atfi.sqrt(1.0 - cosThetaL * cosThetaL)
sinTheta2K = (1.0 - cosThetaK * cosThetaK)
sinTheta2L = (1.0 - cosThetaL * cosThetaL)
sin2ThetaK = (2.0 * sinThetaK * cosThetaK)
cos2ThetaL = (2.0 * cosThetaL * cosThetaL - 1.0)
# Decay density
pdf = (3.0 / 4.0) * (1.0 - FL) * sinTheta2K
pdf += FL * cosThetaK * cosThetaK
pdf += (1.0 / 4.0) * (1.0 - FL) * sin2ThetaK * cos2ThetaL
pdf += (-1.0) * FL * cosThetaK * cosThetaK * cos2ThetaL
pdf += (1.0 / 2.0) * (
1.0 - FL) * AT2 * sinTheta2K * sinTheta2L * atfi.cos(2.0 * phi)
pdf += S5 * sin2ThetaK * sinThetaL * atfi.cos(phi)
return pdf
def invariant_mass_jacobian(self, sample):
"""
sample: [mAB^2, mBC^2]
Return the jacobian determinant (|J|) of tranformation from dmAB^2*dmBC^2 -> |J|*dmAB*dmBC. |J| = 4*mAB*mBC
"""
return (atfi.const(4.0) * atfi.sqrt(self.m2ab(sample)) *
atfi.sqrt(self.m2bc(sample)))
def cos_helicity_angle_dalitz(m2ab, m2bc, md, ma, mb, mc):
"""Calculate cos(helicity angle) for set of two Dalitz plot variables
m2ab, m2bc : Dalitz plot variables (inv. masses squared of AB and BC combinations)
md : mass of the decaying particle
ma, mb, mc : masses of final state particles
:param m2ab:
:param m2bc:
:param md:
:param ma:
:param mb:
:param mc:
"""
md2 = md**2
ma2 = ma**2
mb2 = mb**2
mc2 = mc**2
m2ac = md2 + ma2 + mb2 + mc2 - m2ab - m2bc
mab = atfi.sqrt(m2ab)
mac = atfi.sqrt(m2ac)
mbc = atfi.sqrt(m2bc)
p2a = 0.25 / md2 * (md2 - (mbc + ma)**2) * (md2 - (mbc - ma)**2)
p2b = 0.25 / md2 * (md2 - (mac + mb)**2) * (md2 - (mac - mb)**2)
p2c = 0.25 / md2 * (md2 - (mab + mc)**2) * (md2 - (mab - mc)**2)
eb = (m2ab - ma2 + mb2) / 2. / mab
ec = (md2 - m2ab - mc2) / 2. / mab
pb = atfi.sqrt(eb**2 - mb2)
pc = atfi.sqrt(ec**2 - mc2)
e2sum = (eb + ec)**2
m2bc_max = e2sum - (pb - pc)**2
m2bc_min = e2sum - (pb + pc)**2
return (m2bc_max + m2bc_min - 2. * m2bc) / (m2bc_max - m2bc_min)
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 square_dalitz_plot_jacobian(self, sample):
"""
sample: [mAB^2, mBC^2]
Return the jacobian determinant (|J|) of tranformation from dmAB^2*dmBC^2 -> |J|*dMpr*dThpr where Mpr, Thpr are defined in (AC) frame.
"""
mPrime = self.m_prime_ac(sample)
thPrime = self.theta_prime_ac(sample)
diff_AC = tf.cast(
atfi.sqrt(self.maxac) - atfi.sqrt(self.minac), atfi.fptype())
mAC = atfi.const(0.5) * diff_AC * (
Const(1.0) + atfi.cos(atfi.pi() * mPrime)) + tf.cast(
atfi.sqrt(self.minac), atfi.fptype())
mACSq = mAC * mAC
eAcmsAC = (atfi.const(0.5) *
(mACSq - tf.cast(self.mc2, atfi.fptype()) +
tf.cast(self.ma2, atfi.fptype())) / mAC)
eBcmsAC = (atfi.const(0.5) *
(tf.cast(self.md, atfi.fptype())**2.0 - mACSq -
tf.cast(self.mb2, atfi.fptype())) / mAC)
pAcmsAC = atfi.sqrt(eAcmsAC**2.0 - tf.cast(self.ma2, atfi.fptype()))
pBcmsAC = atfi.sqrt(eBcmsAC**2.0 - tf.cast(self.mb2, atfi.fptype()))
deriv1 = Pi() * atfi.const(0.5) * diff_AC * atfi.sin(
atfi.pi() * mPrime)
deriv2 = Pi() * atfi.sin(atfi.pi() * thPrime)
return atfi.const(4.0) * pAcmsAC * pBcmsAC * mAC * deriv1 * deriv2
def inside(self, x):
"""
Check if the point x=(m2ab, m2bc, cos_theta_a, phi_a, phi_bc) is inside the phase space
"""
m2ab = self.m2ab(x)
m2bc = self.m2bc(x)
mab = atfi.sqrt(m2ab)
costhetaa = self.cos_theta_a(x)
phia = self.phi_a(x)
phibc = self.phi_bc(x)
inside = tf.logical_and(
tf.logical_and(tf.greater(m2ab, self.minab),
tf.less(m2ab, self.maxab)),
tf.logical_and(tf.greater(m2bc, self.minbc),
tf.less(m2bc, self.maxbc)))
if self.macrange:
m2ac = self.msqsum - m2ab - m2bc
inside = tf.logical_and(
inside,
tf.logical_and(tf.greater(m2ac, self.macrange[0]**2),
tf.less(m2ac, self.macrange[1]**2)))
if self.symmetric:
inside = tf.logical_and(inside, tf.greater(m2bc, m2ab))
eb = (m2ab - self.ma2 + self.mb2) / 2. / mab
ec = (self.md2 - m2ab - self.mc2) / 2. / mab
p2b = eb**2 - self.mb2
p2c = ec**2 - self.mc2
inside = tf.logical_and(
inside, tf.logical_and(tf.greater(p2c, 0), tf.greater(p2b, 0)))
pb = atfi.sqrt(p2b)
pc = atfi.sqrt(p2c)
e2bc = (eb + ec)**2
m2bc_max = e2bc - (pb - pc)**2
m2bc_min = e2bc - (pb + pc)**2
inside_phsp = tf.logical_and(
inside,
tf.logical_and(tf.greater(m2bc, m2bc_min), tf.less(m2bc,
m2bc_max)))
inside_theta = tf.logical_and(tf.greater(costhetaa, -1.),
tf.less(costhetaa, 1.))
inside_phi = tf.logical_and(
tf.logical_and(tf.greater(phia, -1. * math.pi),
tf.less(phia, math.pi)),
tf.logical_and(tf.greater(phibc, -1. * math.pi),
tf.less(phibc, math.pi)))
inside_ang = tf.logical_and(inside_theta, inside_phi)
return tf.logical_and(inside_phsp, inside_ang)
def _with_form_factor(self, dataset: dict) -> tf.Tensor:
inv_mass = atfi.sqrt(dataset[self._dynamics_props.inv_mass_name])
m_a = atfi.sqrt(dataset[self._dynamics_props.inv_mass_name_prod1])
m_b = atfi.sqrt(dataset[self._dynamics_props.inv_mass_name_prod2])
meson_radius = self._dynamics_props.meson_radius
l_orbit = self._dynamics_props.orbit_angular_momentum
q_squared = two_body_momentum_squared(inv_mass, m_a, m_b)
return atfi.complex(
atfi.sqrt(
blatt_weisskopf_ff_squared(q_squared, meson_radius, l_orbit)),
atfi.const(0.0),
)
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 polynomial_nonresonant_lineshape(m2,
m0,
coeffs,
ma,
mb,
mc,
md,
lr,
ld,
barrierFactor=True):
"""
2nd order polynomial nonresonant amplitude with orbital barriers
coeffs: list of atfi.complex polynomial coefficients [a0, a1, a2]
"""
def poly(x, cs):
return cs[0] + cs[1] * atfi.complex(
x, atfi.const(0.)) + cs[2] * atfi.complex(x**2, atfi.const(0.))
if barrierFactor:
m = atfi.sqrt(m2)
q = atfk.two_body_momentum(md, m, mc)
q0 = atfk.two_body_momentum(md, m0, mc)
p = atfk.two_body_momentum(m, ma, mb)
p0 = atfk.two_body_momentum(m0, ma, mb)
b1 = orbital_barrier_factor(p, p0, lr)
b2 = orbital_barrier_factor(q, q0, ld)
return poly(m - m0, coeffs) * atfi.complex(b1 * b2, atfi.const(0.))
else:
return poly(m - m0, coeffs)
def exponential_nonresonant_lineshape(m2,
m0,
alpha,
ma,
mb,
mc,
md,
lr,
ld,
barrierFactor=True):
"""
Exponential nonresonant amplitude with orbital barriers
"""
if barrierFactor:
m = atfi.sqrt(m2)
q = atfk.two_body_momentum(md, m, mc)
q0 = atfk.two_body_momentum(md, m0, mc)
p = atfk.two_body_momentum(m, ma, mb)
p0 = atfk.two_body_momentum(m0, ma, mb)
b1 = orbital_barrier_factor(p, p0, lr)
b2 = orbital_barrier_factor(q, q0, ld)
return atfi.complex(b1 * b2 * atfi.exp(-alpha * (m2 - m0**2)),
atfi.const(0.))
else:
return atfi.complex(atfi.exp(-alpha * (m2 - m0**2)), atfi.const(0.))
def subthreshold_breit_wigner_lineshape(m2,
m0,
gamma0,
ma,
mb,
mc,
md,
dr,
dd,
lr,
ld,
barrier_factor=True):
"""
Breit-Wigner amplitude (with the mass under kinematic threshold)
with Blatt-Weisskopf formfactors, mass-dependent width and orbital barriers
"""
m = atfi.sqrt(m2)
mmin = ma + mb
mmax = md - mc
tanhterm = atfi.tanh((m0 - ((mmin + mmax) / 2.)) / (mmax - mmin))
m0eff = mmin + (mmax - mmin) * (1. + tanhterm) / 2.
q = atfk.two_body_momentum(md, m, mc)
q0 = atfk.two_body_momentum(md, m0eff, mc)
p = atfk.two_body_momentum(m, ma, mb)
p0 = atfk.two_body_momentum(m0eff, ma, mb)
ffr = blatt_weisskopf_ff(p, p0, dr, lr)
ffd = blatt_weisskopf_ff(q, q0, dd, ld)
width = mass_dependent_width(m, m0, gamma0, p, p0, ffr, lr)
bw = relativistic_breit_wigner(m2, m0, width)
ff = ffr * ffd
if barrier_factor:
b1 = orbital_barrier_factor(p, p0, lr)
b2 = orbital_barrier_factor(q, q0, ld)
ff *= b1 * b2
return bw * atfi.complex(ff, atfi.const(0.))
def breit_wigner_lineshape(m2,
m0,
gamma0,
ma,
mb,
mc,
md,
dr,
dd,
lr,
ld,
barrier_factor=True,
ma0=None,
md0=None):
"""
Breit-Wigner amplitude with Blatt-Weisskopf formfactors, mass-dependent width and orbital barriers
"""
m = atfi.sqrt(m2)
q = atfk.two_body_momentum(md, m, mc)
q0 = atfk.two_body_momentum(md if md0 is None else md0, m0, mc)
p = atfk.two_body_momentum(m, ma, mb)
p0 = atfk.two_body_momentum(m0, ma if ma0 is None else ma0, mb)
ffr = blatt_weisskopf_ff(p, p0, dr, lr)
ffd = blatt_weisskopf_ff(q, q0, dd, ld)
width = mass_dependent_width(m, m0, gamma0, p, p0, ffr, lr)
bw = relativistic_breit_wigner(m2, m0, width)
ff = ffr * ffd
if barrier_factor:
b1 = orbital_barrier_factor(p, p0, lr)
b2 = orbital_barrier_factor(q, q0, ld)
ff *= b1 * b2
return bw * atfi.complex(ff, atfi.const(0.))
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 norm(vec):
"""Calculate norm of 3-vector
:param vec:
"""
return atfi.sqrt(tf.reduce_sum(vec * vec, -1))
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 pt(vector):
"""Return transverse (X-Y) component of the input Lorentz or 3-vector
:param vector: input vector (Lorentz or 3-vector)
:returns: vector of transverse components
"""
return atfi.sqrt(x_component(vector)**2 + y_component(vector)**2)
def mass(vector):
"""Calculate mass scalar for Lorentz 4-momentum
vector : input Lorentz momentum vector
:param vector:
"""
return atfi.sqrt(tf.reduce_sum(vector * vector * metric_tensor(), -1))
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 two_body_momentum(md, ma, mb):
"""Momentum of two-body decay products D->AB in the D rest frame
:param md:
:param ma:
:param mb:
"""
return atfi.sqrt(two_body_momentum_squared(md, ma, mb))
def mass(vector):
"""
Calculate mass scalar for Lorentz 4-momentum vector
:param vector: input Lorentz momentum vector
:returns: scalar invariant mass
"""
return atfi.sqrt(mass_squared(vector))
def norm(vec):
"""
Calculate norm of 3-vector
:param vec: Input 3-vector
:returns: Scalar norm
"""
return atfi.sqrt(tf.reduce_sum(vec * vec, -1))
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 nonresonant_lass_lineshape(m2ab, a, r, ma, mb):
"""
LASS line shape, nonresonant part
"""
m = atfi.sqrt(m2ab)
q = atfk.two_body_momentum(m, ma, mb)
cot_deltab = 1. / a / q + 1. / 2. * r * q
ampl = atfi.cast_complex(m) / atfi.complex(q * cot_deltab, -q)
return ampl