def inside(self, x):
inside = self.simple_inside(x)
if self.mab1range or self.mab2range:
(pa1, pa2, pb1, pb2) = self.final_state_momenta(x)
if self.mab1range:
mab1 = atfk.mass(pa1 + pa2 + pb1)
inside = tf.logical_and(
inside,
tf.logical_and(
tf.greater(mab1, self.mab1range[0]),
tf.less(mab1, self.mab1range[1]),
),
)
if self.mab2range:
mab2 = atfk.mass(pa1 + pa2 + pb2)
inside = tf.logical_and(
inside,
tf.logical_and(
tf.greater(mab2, self.mab2range[0]),
tf.less(mab2, self.mab2range[1]),
),
)
return inside
def selection(sample, cuts, rnd, constant_cuts=False):
ptcut = cuts[0]
pcut = cuts[1]
d_ptcut = cuts[2]
max_ptcut = cuts[3]
sum_ptcut = cuts[4]
dalitz_sample = dalitz_phase_space.from_square_dalitz_plot(
sample[:, 0], sample[:, 1])
# Random momenta for a given Dalitz plot sample
mom = dalitz_phase_space.final_state_momenta(
dalitz_phase_space.m2ac(
dalitz_sample
), # Note AB<->AC because FroSquareDalitzPlot works with AC
dalitz_phase_space.m2bc(dalitz_sample)
) # Generate momenta in a frame of the decaying particle
mom = generate_rotation_and_boost(
mom, md, meanpt, d_ptcut,
rnd[:, :]) # Rotate and boost to the lab frame
# Apply cuts
sel = atfi.greater(atfk.pt(mom[0]), ptcut)
sel = atfi.logical_and(sel, atfi.greater(atfk.pt(mom[1]), ptcut))
sel = atfi.logical_and(sel, atfi.greater(atfk.pt(mom[2]), ptcut))
sel = atfi.logical_and(
sel, atfi.greater(atfk.pt(mom[0] + mom[1] + mom[2]), d_ptcut))
sel = atfi.logical_and(sel, atfi.greater(atfk.p(mom[0]), pcut))
sel = atfi.logical_and(sel, atfi.greater(atfk.p(mom[1]), pcut))
sel = atfi.logical_and(sel, atfi.greater(atfk.p(mom[2]), pcut))
sel = atfi.logical_and(
sel,
atfi.greater(
atfi.max(atfi.max(atfk.pt(mom[0]), atfk.pt(mom[1])),
atfk.pt(mom[2])), max_ptcut))
sel = atfi.logical_and(
sel,
atfi.greater(
atfk.pt(mom[0]) + atfk.pt(mom[1]) + atfk.pt(mom[2]), sum_ptcut))
sel = atfi.logical_and(sel, atfi.greater(atfk.eta(mom[0]), etamin))
sel = atfi.logical_and(sel, atfi.greater(atfk.eta(mom[1]), etamin))
sel = atfi.logical_and(sel, atfi.greater(atfk.eta(mom[2]), etamin))
sel = atfi.logical_and(sel, atfi.less(atfk.eta(mom[0]), etamax))
sel = atfi.logical_and(sel, atfi.less(atfk.eta(mom[1]), etamax))
sel = atfi.logical_and(sel, atfi.less(atfk.eta(mom[2]), etamax))
m2ab = atfk.mass(mom[0] + mom[1])**2
m2bc = atfk.mass(mom[1] + mom[2])**2
mprime = sample[:, 0]
thetaprime = sample[:, 1]
arrays = []
outlist = [mprime, thetaprime]
if not constant_cuts:
outlist += [ptcut, pcut, d_ptcut, max_ptcut, sum_ptcut]
for i in outlist:
arrays += [i[sel]]
return arrays
def get_phasespace_vars(PLb_lab, PLc_lab, PLepton_lab):
PW_lab = PLb_lab - PLc_lab
PLb_Wlab = atfk.boost_to_rest(PLb_lab, PW_lab)
PLepton_Wlab = atfk.boost_to_rest(PLepton_lab, PW_lab)
q2 = atfk.mass(PW_lab)**2 / 1e6 #Convert to GeV^2
costhl = atfk.scalar_product(
atfk.unit_vector(atfk.spatial_components(PLepton_Wlab)),
-1. * atfk.unit_vector(atfk.spatial_components(PLb_Wlab)))
return q2, costhl
def kinematic_fit(mfit, moms):
"""
Kinematic fit to a fixed invariant mass for a multibody decay.
Returns the fitted mass and the list of final state 4-momenta.
"""
mcorr = atfk.mass(sum(moms))
for l in range(3):
dm2 = mcorr**2 - mfit**2
delta = momentum_scale(dm2, moms)
moms2 = []
for mom in moms:
m2 = atfk.mass(mom)**2
momvec = atfk.spatial_components(mom) * atfk.scalar(
1 + delta * momentum_resolution(atfk.p(mom)))
mom2 = atfk.lorentz_vector(momvec,
atfi.sqrt(m2 + atfk.norm(momvec)**2))
moms2 += [mom2]
moms = moms2
mcorr = atfk.mass(sum(moms))
return mcorr, moms
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))
def generate_selection(cuts, rnd, constant_cuts=False):
"""
Call generation of fully combinatorial or combinatorial with intermediate K* or rho resonances with specified fractions.
Apply cuts to the final state particles and fill in output arrays.
"""
meankpt = cuts[0]
meanpipt = cuts[1]
ptcut = cuts[2]
pcut = cuts[3]
meankstarpt = cuts[4]
meanrhopt = cuts[5]
kstarfrac = cuts[6]
rhofrac = cuts[7]
p4k_1, p4pi1_1, p4pi2_1 = generate_combinatorial(cuts, rnd)
p4k_2, p4pi1_2, p4pi2_2 = generate_kstar(cuts, rnd)
p4k_3, p4pi1_3, p4pi2_3 = generate_rho(cuts, rnd)
thr1 = 1. - kstarfrac - rhofrac
thr2 = 1. - rhofrac
cond1 = atfi.stack(4 * [rnd[:, 11] < thr1], axis=1)
cond2 = atfi.stack(4 * [rnd[:, 11] < thr2], axis=1)
p4k = tf.where(cond1, p4k_1, tf.where(cond2, p4k_2, p4k_3))
p4pi1 = tf.where(cond1, p4pi1_1, tf.where(cond2, p4pi1_2, p4pi1_3))
p4pi2 = tf.where(cond1, p4pi2_1, tf.where(cond2, p4pi2_2, p4pi2_3))
mb = atfk.mass(p4k + p4pi1 + p4pi2)
mfit, moms = kinematic_fit(atfi.const(md), [p4k, p4pi1, p4pi2])
sel = tf.greater(atfk.p(moms[0]), pcut)
sel = tf.logical_and(sel, tf.greater(atfk.p(moms[1]), pcut))
sel = tf.logical_and(sel, tf.greater(atfk.p(moms[2]), pcut))
sel = tf.logical_and(sel, tf.greater(atfk.pt(moms[0]), ptcut))
sel = tf.logical_and(sel, tf.greater(atfk.pt(moms[1]), ptcut))
sel = tf.logical_and(sel, tf.greater(atfk.pt(moms[2]), ptcut))
m2kpi = atfk.mass(moms[0] + moms[1])**2
m2pipi = atfk.mass(moms[1] + moms[2])**2
sample = dlz_phsp.from_vectors(m2kpi, m2pipi)
mprime = dlz_phsp.m_prime_bc(sample)
thetaprime = dlz_phsp.theta_prime_bc(sample)
sel = tf.logical_and(
sel,
observables_phase_space.inside(
tf.stack([mprime, thetaprime, mb], axis=1)))
observables = []
outlist = [mprime, thetaprime, mb, m2kpi, m2pipi]
if not constant_cuts:
outlist += [
meankpt, meanpipt, ptcut, pcut, meankstarpt, meanrhopt, kstarfrac,
rhofrac
]
for i in outlist:
observables += [tf.boolean_mask(i, sel)]
return observables