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 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 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 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 helicity_amplitude_3body(thetaR, phiR, thetaA, phiA, spinD, spinR, mu,
lambdaR, lambdaA, lambdaB, lambdaC):
"""Calculate complex helicity amplitude for the 3-body decay D->ABC
thetaR, phiR : polar and azimuthal angles of AB resonance in D rest frame
thetaA, phiA : polar and azimuthal angles of A in AB rest frame
spinD : D spin
spinR : spin of the intermediate R resonance
mu : D spin projection onto z axis
lambdaR : R resonance helicity
lambdaA : A helicity
lambdaB : B helicity
lambdaC : C helicity
:param thetaR:
:param phiR:
:param thetaA:
:param phiA:
:param spinD:
:param spinR:
:param mu:
:param lambdaR:
:param lambdaA:
:param lambdaB:
:param lambdaC:
"""
lambda1 = lambdaR - lambdaC
lambda2 = lambdaA - lambdaB
ph = (mu - lambda1) / 2. * phiR + (lambdaR - lambda2) / 2. * phiA
d_terms = wigner_small_d(thetaR, spinD, mu, lambda1) * \
wigner_small_d(thetaA, spinR, lambdaR, lambda2)
h = atfi.complex(d_terms * atfi.cos(ph), d_terms * atfi.sin(ph))
return h
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 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 rotate_euler(v, phi, theta, psi):
"""Perform 3D rotation of the 3-vector
v : vector to be rotated
phi, theta, psi : Euler angles in Z-Y-Z convention
:param v:
:param phi:
:param theta:
:param psi:
"""
# rotate Z (phi)
c1 = atfi.cos(phi)
s1 = atfi.sin(phi)
c2 = atfi.cos(theta)
s2 = atfi.sin(theta)
c3 = atfi.cos(psi)
s3 = atfi.sin(psi)
# rotate Y (theta)
fzx2 = -s2 * c1
fzy2 = s2 * s1
fzz2 = c2
# rotate Z (psi)
fxx3 = c3 * c2 * c1 - s3 * s1
fxy3 = -c3 * c2 * s1 - s3 * c1
fxz3 = c3 * s2
fyx3 = s3 * c2 * c1 + c3 * s1
fyy3 = -s3 * c2 * s1 + c3 * c1
fyz3 = s3 * s2
# Transform v
vx = x_component(v)
vy = y_component(v)
vz = z_component(v)
_vx = fxx3 * vx + fxy3 * vy + fxz3 * vz
_vy = fyx3 * vx + fyy3 * vy + fyz3 * vz
_vz = fzx2 * vx + fzy2 * vy + fzz2 * vz
return vector(_vx, _vy, _vz)
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 resonant_lass_lineshape(m2ab, m0, gamma0, a, r, ma, mb):
"""
LASS line shape, resonant part
"""
m = atfi.sqrt(m2ab)
q0 = atfk.two_body_momentum(m0, ma, mb)
q = atfk.two_body_momentum(m, ma, mb)
cot_deltab = 1. / a / q + 1. / 2. * r * q
phase = atfi.atan(1. / cot_deltab)
width = gamma0 * q / m * m0 / q0
ampl = relativistic_breit_wigner(m2ab, m0, width) * atfi.complex(
atfi.cos(phase), atfi.sin(phase)) * atfi.cast_complex(
m2ab * gamma0 / q0)
return ampl
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 normal_random(rnd1, rnd2):
"""
Normal distribution from two random numbers
"""
return atfi.sqrt(-2. * atfi.log(rnd1)) * atfi.cos(2. * math.pi * rnd2)
