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 momentum_scale(dm, moms):
"""
Function to calculate the momentum scale factor for kinematic fit
dm : invariant mass shift from the desired value
moms : list of 4-momenta of the final state particles
"""
psum = sum(moms) # sum of 4-momenta
pvecsum = atfk.spatial_components(psum)
esum = atfk.time_component(psum)
dedd = atfi.const(0.)
pdpdd = atfi.const(0.)
for mom in moms:
pm = atfk.p(mom) # Absolute momentum
em = atfk.time_component(mom) # Energy
pvec = atfk.spatial_components(mom) # 3-momentum
s = momentum_resolution(pm)
dedd += s * pm**2 / em
pdpdd += atfk.scalar_product(pvecsum, pvec) * s
return -dm / (2. * esum * dedd - 2. * pdpdd)
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
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
moms = final_state_momenta(toy_data)
# Apply random rotation and boost to lab frame
rnd = tf.random.uniform((len(toy_data), 6), dtype=tf.float64)
boosted_moms = random_rotation_and_boost(moms, rnd)
print(f"Number of events generated: {len(boosted_moms[0])}")
# Filter data according to "trigger conditions"
# (minimun p and pT for tracks, and minimum angles between pairs of tracks).
filtered_data = toy_data[
(atfk.pt(boosted_moms[0]) > 1.0)
& (atfk.pt(boosted_moms[1]) > 1.0)
& (atfk.pt(boosted_moms[2]) > 0.5)
& (atfk.pt(boosted_moms[3]) > 0.5)
& (atfk.p(boosted_moms[0]) > 5.0)
& (atfk.p(boosted_moms[1]) > 5.0)
& (atfk.p(boosted_moms[2]) > 3.0)
& (atfk.p(boosted_moms[3]) > 3.0)
& (
atfk.scalar_product(
atfk.unit_vector(boosted_moms[0]), atfk.unit_vector(boosted_moms[1])
)
< 0.9999
)
& (
atfk.scalar_product(
atfk.unit_vector(boosted_moms[0]), atfk.unit_vector(boosted_moms[2])
)
< 0.9999
)