def orbital_barrier_factor(p, p0, l):
"""
Orbital barrier factor
"""
if l == 0: return atfi.ones(p)
if l == 1: return (p / p0)
if l >= 2: return (p / p0)**l
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 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 _model(a1r, a1i, a2r, a2i, a3r, a3i, switches=4 * [1]):
a1 = atfi.complex(a1r, a1i)
a2 = atfi.complex(a2r, a2i)
a3 = atfi.complex(a3r, a3i)
ampl = atfi.cast_complex(atfi.ones(m2ab)) * atfi.complex(
atfi.const(0.0), atfi.const(0.0))
if switches[0]:
ampl += a1 * bw1 * hel_ab
if switches[1]:
ampl += a2 * bw2 * hel_bc
if switches[2]:
ampl += a3 * bw3 * hel_ac
if switches[3]:
ampl += atfi.cast_complex(atfi.ones(m2ab)) * atfi.complex(
atfi.const(5.0), atfi.const(0.0))
return atfd.density(ampl)
def axes_before_rotation(pb):
"""Calculate old (before rotation) axes in the frame aligned with the momentum vector pb
:param pb:
"""
z1 = unit_vector(
spatial_components(pb)) # New z-axis is in the direction of pb
eb = time_component(pb)
z0 = vector(atfi.zeros(eb), atfi.zeros(eb),
atfi.ones(eb)) # Old z-axis vector
x0 = vector(atfi.ones(eb), atfi.zeros(eb),
atfi.zeros(eb)) # Old x-axis vector
sp = scalar_product(z1, z0)
a0 = z0 - z1 * scalar(sp) # vector in z-pb plane perpendicular to z0
x1 = tf.where(tf.equal(sp, 1.0), x0, -unit_vector(a0))
y1 = vector_product(z1, x1) # New y-axis
x = vector(x_component(x1), x_component(y1), x_component(z1))
y = vector(y_component(x1), y_component(y1), y_component(z1))
z = vector(z_component(x1), z_component(y1), z_component(z1))
return (x, y, z)
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 model(x, mrho, wrho, mkst, wkst, a1r, a1i, a2r, a2i, a3r, a3i, switches):
a1 = atfi.complex(a1r, a1i)
a2 = atfi.complex(a2r, a2i)
a3 = atfi.complex(a3r, a3i)
m2ab = phsp.m2ab(x)
m2bc = phsp.m2bc(x)
m2ac = phsp.m2ac(x)
hel_ab = phsp.cos_helicity_ab(x)
hel_bc = phsp.cos_helicity_bc(x)
hel_ac = phsp.cos_helicity_ac(x)
ampl = atfi.complex(atfi.const(0.0), atfi.const(0.0))
if switches[0]:
ampl += (
a1
* atfd.breit_wigner_lineshape(
m2ab, mkst, wkst, mpi, mk, mpi, md, rd, rr, 1, 1
)
* atfd.helicity_amplitude(hel_ab, 1)
)
if switches[1]:
ampl += (
a2
* atfd.breit_wigner_lineshape(
m2bc, mkst, wkst, mpi, mk, mpi, md, rd, rr, 1, 1
)
* atfd.helicity_amplitude(hel_bc, 1)
)
if switches[2]:
ampl += (
a3
* atfd.breit_wigner_lineshape(
m2ac, mrho, wrho, mpi, mpi, mk, md, rd, rr, 1, 1
)
* atfd.helicity_amplitude(hel_ac, 1)
)
if switches[3]:
ampl += atfi.cast_complex(atfi.ones(m2ab)) * atfi.complex(
atfi.const(5.0), atfi.const(0.0)
)
return atfd.density(ampl)
def axes_after_rotation(pb, oldaxes=None):
"""Calculate new (rotated) axes aligned with the momentum vector pb
:param pb:
:param oldaxes: (Default value = None)
"""
z1 = unit_vector(
spatial_components(pb)) # New z-axis is in the direction of pb
eb = time_component(pb)
zeros = atfi.zeros(eb)
ones = atfi.ones(eb)
# Old z-axis vector
z0 = vector(zeros, zeros, ones) if oldaxes == None else oldaxes[2]
# Old x-axis vector
x0 = vector(ones, zeros, zeros) if oldaxes == None else oldaxes[0]
sp = scalar_product(z1, z0)
a0 = z0 - z1 * scalar(sp) # vector in z-pb plane perpendicular to z0
x1 = tf.where(scalar(tf.equal(sp, 1.)), x0, -unit_vector(a0))
y1 = vector_product(z1, x1) # New y-axis
return (x1, y1, z1)
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))