def main() :
nev = int(sys.argv[1])
outfile = sys.argv[2]
chunk_size = 1000000 # Events will be generated in chunks of this size
atfi.set_seed(nev)
n = 0 # Current tuple size
arrays = []
while(True) :
# Create Dalitz plot sample
unfiltered_sample = observables_phase_space.unfiltered_sample(chunk_size) # Unfiltered array
sample = observables_phase_space.filter(unfiltered_sample)
size = sample.shape[0]
print(f"Filtered chunk size = {size}")
# Generate final state momenta from Dalitz plot and run through selection
rnd = tf.random.uniform([size, random_array_size], dtype = atfi.fptype() ) # Auxiliary random array
array = atfi.stack(generate_selection( true_cuts, rnd, constant_cuts = True ), axis = 1)
arrays += [ array ]
# Increment counters and check if we are done
size = array.shape[0]
n += size
if n > nev : break
print(f"Selected size = {n}, last = {size}")
tfr.write_tuple(outfile, atfi.concat(arrays, axis = 0)[:nev,:], observables_toys)
def main():
nev = 1000000
outfile = "toy_tuple.root"
atfi.set_seed(nev + 1)
chunk_size = 1000000 # Events will be generated in chunks of this size
bounds = {i[0]: (i[2], i[3])
for i in parameters_list
} # Bounds and exponential factor for generation of cuts
branches = generated_variables + [i[0] for i in parameters_list]
n = 0 # Current tuple size
arrays = []
while (True):
# Generate final state momenta from Dalitz plot and run through selection
rnd = tf.random.uniform([chunk_size, random_array_size + len(bounds)],
dtype=atfi.fptype()) # Auxiliary random array
array = atfi.stack(generate_candidates_and_cuts(rnd), axis=1)
arrays += [array]
# Increment counters and check if we are done
size = array.shape[0]
n += size
if n > nev: break
print(f"Selected size = {n}, last = {size}")
tfr.write_tuple(outfile, atfi.concat(arrays, axis=0)[:nev, :], branches)
def main():
nev = 2000000
outfile = "toy_tuple.root"
atfi.set_seed(nev + 1)
chunk_size = 1000000 # Events will be generated in chunks of this size
bounds = {i[0]: (i[2], i[3])
for i in parameters_list
} # Bounds and exponential factor for generation of cuts
branches = observables_toys + [i[0] for i in parameters_list]
n = 0 # Current tuple size
arrays = []
while (True):
# Create Dalitz plot sample
unfiltered_sample = observables_phase_space.unfiltered_sample(
chunk_size) # Unfiltered array
sample = observables_phase_space.filter(unfiltered_sample)
size = sample.shape[0]
print(f"Filtered chunk size = {size}")
# Generate final state momenta from Dalitz plot and run through selection
rnd = tf.random.uniform([size, random_array_size],
dtype=atfi.fptype()) # Auxiliary random array
array = atfi.stack(selection_with_random_cuts(sample, rnd), axis=1)
arrays += [array]
# Increment counters and check if we are done
size = array.shape[0]
n += size
if n > nev: break
print(f"Selected size = {n}, last = {size}")
tfr.write_tuple(outfile, atfi.concat(arrays, axis=0)[:nev, :], branches)
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