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 helicity_angles_3body(pa, pb, pc):
"""Calculate 4 helicity angles for the 3-body D->ABC decay defined as:
theta_r, phi_r : polar and azimuthal angles of the AB resonance in the D rest frame
theta_a, phi_a : polar and azimuthal angles of the A in AB rest frame
:param pa:
:param pb:
:param pc:
"""
theta_r = atfi.acos(-z_component(pc) / norm(spatial_components(pc)))
phi_r = atfi.atan2(-y_component(pc), -x_component(pc))
pa_prime = lorentz_vector(
rotate_euler(spatial_components(pa), -phi_r,
atfi.pi() - theta_r, phi_r), time_component(pa))
pb_prime = lorentz_vector(
rotate_euler(spatial_components(pb), -phi_r,
atfi.pi() - theta_r, phi_r), time_component(pb))
w = time_component(pa) + time_component(pb)
pab = lorentz_vector(-(pa_prime + pb_prime) / scalar(w), w)
pa_prime2 = lorentz_boost(pa_prime, pab)
theta_a = atfi.acos(
z_component(pa_prime2) / norm(spatial_components(pa_prime2)))
phi_a = atfi.atan2(y_component(pa_prime2), x_component(pa_prime2))
return (theta_r, phi_r, theta_a, phi_a)
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 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 helicity_angles_4body(pa, pb, pc, pd):
"""Calculate 4 helicity angles for the 4-body E->ABCD decay defined as:
theta_ab, phi_ab : polar and azimuthal angles of the AB resonance in the E rest frame
theta_cd, phi_cd : polar and azimuthal angles of the CD resonance in the E rest frame
theta_ac, phi_ac : polar and azimuthal angles of the AC resonance in the E rest frame
theta_bd, phi_bd : polar and azimuthal angles of the BD resonance in the E rest frame
theta_ad, phi_ad : polar and azimuthal angles of the AD resonance in the E rest frame
theta_bc, phi_bc : polar and azimuthal angles of the BC resonance in the E rest frame
phi_ab_cd : azimuthal angle between AB and CD
phi_ac_bd : azimuthal angle between AC and BD
phi_ad_bc : azimuthal angle between AD and BC
:param pa:
:param pb:
:param pc:
:param pd:
"""
theta_r = atfi.acos(-z_component(pc) / norm(spatial_components(pc)))
phi_r = atfi.atan2(-y_component(pc), -x_component(pc))
pa_prime = lorentz_vector(
rotate_euler(spatial_components(pa), -phi_r,
atfi.pi() - theta_r, phi_r),
time_component(pa),
)
pb_prime = lorentz_vector(
rotate_euler(spatial_components(pb), -phi_r,
atfi.pi() - theta_r, phi_r),
time_component(pb),
)
w = time_component(pa) + time_component(pb)
pab = lorentz_vector(-(pa_prime + pb_prime) / scalar(w), w)
pa_prime2 = lorentz_boost(pa_prime, pab)
theta_a = atfi.acos(
z_component(pa_prime2) / norm(spatial_components(pa_prime2)))
phi_a = atfi.atan2(y_component(pa_prime2), x_component(pa_prime2))
return (theta_r, phi_r, theta_a, phi_a)
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
sys.path.append("../")
#os.environ["CUDA_VISIBLE_DEVICES"] = "" # Do not use GPU
import amplitf.interface as atfi
import amplitf.kinematics as atfk
import amplitf.dynamics as atfd
import amplitf.toymc as atft
from amplitf.phasespace.rectangular_phasespace import RectangularPhaseSpace
from ROOT import TFile
if __name__ == "__main__":
# Four body angular phase space is described by 3 angles.
phsp = RectangularPhaseSpace(
((-1., 1.), (-1., 1.), (-atfi.pi(), atfi.pi())))
# Fit parameters of the model
FL = 0.770
AT2 = 0.200
S5 = -0.100
### Start of model description
def model(x):
# Get phase space variables
cosThetaK = phsp.coordinate(x, 0)
cosThetaL = phsp.coordinate(x, 1)
phi = phsp.coordinate(x, 2)
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))
import amplitf.likelihood as atfl
import amplitf.optimisation as atfo
from amplitf.phasespace.rectangular_phasespace import RectangularPhaseSpace
# Import TFA modules
import tfa.toymc as tft
from ROOT import TFile
tf.config.experimental_run_functions_eagerly(False)
if __name__ == "__main__":
# Four body angular phase space is described by 3 angles.
phsp = RectangularPhaseSpace(
((-1.0, 1.0), (-1.0, 1.0), (-atfi.pi(), atfi.pi())))
# Fit parameters of the model
FL = atfo.FitParameter("FL", 0.770, 0.000, 1.000, 0.01)
AT2 = atfo.FitParameter("AT2", 0.200, -1.000, 1.000, 0.01)
S5 = atfo.FitParameter("S5", -0.100, -1.000, 1.000, 0.01)
pars = [FL, AT2, S5]
### Start of model description
@atfi.function
def model(x):
# Get phase space variables
cosThetaK = phsp.coordinate(x, 0)
cosThetaL = phsp.coordinate(x, 1)