def pt_cut_3of3(pa, pb, p1, p2, p3):
""" Ensures that both p1 and p2 pass the pt_cut
returns a boolean mask and the list of true indices
"""
# Put all pt on a stack
all_p = tf.stack([p1, p2, p3])
# First check that all 3 pass the pt cut
all_pt2 = pt2many(all_p)
ptpass = tf.reduce_all(all_pt2 > pt2_cut, axis=0)
# Now select the two hardest jets and apply the rest of the cuts on those two
pts, idx = tf.math.top_k(tf.transpose(all_pt2), k=2, sorted=False)
pjsT = tf.gather(tf.transpose(all_p, perm=[2, 0, 1]), idx, batch_dims=1)
pjs = tf.transpose(pjsT, perm=[1, 2, 0])
# Check that the two hardest jets pass the rapidity delta cut
ypass = deltay_cut(pjs[0], pjs[1])
# and the mjj cut
mjjpass = mjj_cut(pjs[0], pjs[1])
# Check that all jets pass the rdistance check
rpass = rcone_cut(p1, p2, p3)
# Ensure that all invariants are above some threshold
tech_cut_pass = invariant_cut(pa, pb, p1, p2, p3)
jetpass = tf.reduce_all([rpass, tech_cut_pass, ypass, mjjpass], axis=0)
stripe, idx = _condition_to_idx(ptpass, jetpass)
return stripe, idx, pts[:, 0]
def alphas_first_subgrid(
shape,
a_q2,
log_q2min,
log_q2max,
padded_q2,
s_q2,
actual_padded,
):
"""
First subgrid interpolation
Calls
alphas_interpolate (basic interpolation) (0)
Parameters
----------
shape: tf.tensor of shape [None]
final output shape to scatter points into
For other parameters refer to subgrid.py:alphas_interpolate
Returns
----------
tf.tensor of shape `shape`
alphas interpolated values for each query point
"""
res = tf.zeros(shape, dtype=DTYPE)
# --------------------------------------------------------------------
# normal interpolation
stripe, f_idx = _condition_to_idx(a_q2 >= log_q2min, a_q2 < log_q2max)
if tf.math.equal(f_idx, 0) is not None:
in_q2 = tf.boolean_mask(a_q2, stripe)
ff_f = alphas_interpolate(in_q2, padded_q2, s_q2, actual_padded)
res = tf.tensor_scatter_nd_update(res, f_idx, ff_f)
# --------------------------------------------------------------------
# lowq2
stripe = a_q2 < log_q2min
f_idx = int_me(tf.where(stripe))
if tf.math.equal(f_idx, 0) is not None:
in_q2 = tf.boolean_mask(a_q2, stripe)
m = tf.math.log(actual_padded[2]/actual_padded[1])\
/(padded_q2[2] - padded_q2[1])
ff_f = actual_padded[1] * tf.math.pow(
tf.math.exp(in_q2) / tf.math.exp(padded_q2[1]), m)
res = tf.tensor_scatter_nd_update(res, f_idx, ff_f)
return res
def pt_cut_2of2(p1, p2):
""" Ensures that both p1 and p2 pass the pt_cut
returns a boolean mask and the list of true indices
"""
pt21 = pt2(p1)
pt22 = pt2(p2)
p1pass = pt21 > pt2_cut
p2pass = pt22 > pt2_cut
deltaypass = deltay_cut(p1, p2)
mjjpass = mjj_cut(p1, p2)
jetpass = tf.reduce_all([p1pass, p2pass, deltaypass], axis=0)
stripe, idx = _condition_to_idx(jetpass, mjjpass)
return stripe, idx, tf.reduce_max([pt21, pt22], axis=0)
def alphas_last_subgrid(
shape,
a_q2,
log_q2min,
log_q2max,
padded_q2,
s_q2,
actual_padded,
):
"""
Last subgrid interpolation.
Calls
alphas_interpolate: (0)
Parameters
----------
shape: tf.tensor of shape [None]
final output shape to scatter points into
For other parameters see :py:func:`pdfflow.alphas_region_interpolator.alphas_interpolate`
Returns
----------
tf.tensor, shape: `shape`
alphas interpolated values for each query point
"""
# Generate all conditions for all stripes
res = tf.zeros(shape, dtype=DTYPE)
# --------------------------------------------------------------------
# normal interpolation
stripe, f_idx = _condition_to_idx(a_q2 >= log_q2min, a_q2 <= log_q2max)
if tf.math.equal(f_idx, 0) is not None:
# Check whether there are any points in this region
# if there are, execute normal_interpolation
in_q2 = tf.boolean_mask(a_q2, stripe)
ff_f = alphas_interpolate(in_q2, padded_q2, s_q2, actual_padded)
res = tf.tensor_scatter_nd_update(res, f_idx, ff_f)
# --------------------------------------------------------------------
# high q2
stripe = a_q2 > log_q2max
f_idx = int_me(tf.where(stripe))
if tf.math.equal(f_idx, 0) is not None:
ff_f = tf.ones_like(f_idx[:, 0], dtype=DTYPE) * actual_padded[-2]
res = tf.tensor_scatter_nd_update(res, f_idx, ff_f)
return res
def alphas_inner_subgrid(
shape,
a_q2,
log_q2min,
log_q2max,
padded_q2,
s_q2,
actual_padded,
):
"""
Inner (non-first and non-last) alphas subgrid interpolation
Calls
alphas_interpolate (basic interpolation) (0)
Parameters
----------
shape: tf.tensor of shape [None]
final output shape to scatter points into
For other parameters refer to subgrid.py:alphas_interpolate
Returns
----------
tf.tensor of shape `shape`
alphas interpolated values for each query point
"""
res = tf.zeros(shape, dtype=DTYPE)
# --------------------------------------------------------------------
# normal interpolation
stripe, f_idx = _condition_to_idx(a_q2 >= log_q2min, a_q2 < log_q2max)
if tf.math.equal(f_idx, 0) is not None:
in_q2 = tf.boolean_mask(a_q2, stripe)
ff_f = alphas_interpolate(in_q2, padded_q2, s_q2, actual_padded)
res = tf.tensor_scatter_nd_update(res, f_idx, ff_f)
return res