def calculate_selection(self, syst_tag, events):
"""
"""
electrons = events.ele
electrons["label"] = -1 * awkward.ones_like(electrons.pt)
if not self.is_data:
electrons["label"] = awkward.where(
electrons.genPartFlav == 1, awkward.ones_like(electrons.label),
electrons.label)
electrons["label"] = awkward.where(
(electrons.genPartFlav == 3) | (electrons.genPartFlav == 4) |
(electrons.genPartFlav == 5),
awkward.zeros_like(electrons.label), electrons.label)
fields = [x for x in events.fields if x not in ["ele", "Electron"]]
for x in fields:
if x == "ZCand":
electrons[x] = awkward.firsts(events[x])
else:
electrons[x] = events[x]
electrons = awkward.flatten(electrons)
dummy_cut = electrons.pt >= 0
return dummy_cut, electrons
def apply_roccor(df, rochester, is_mc):
if is_mc:
hasgen = ~np.isnan(ak.fill_none(df.Muon.matched_gen.pt, np.nan))
mc_rand = np.random.rand(*ak.to_numpy(ak.flatten(df.Muon.pt)).shape)
mc_rand = ak.unflatten(mc_rand, ak.num(df.Muon.pt, axis=1))
corrections = np.array(ak.flatten(ak.ones_like(df.Muon.pt)))
errors = np.array(ak.flatten(ak.ones_like(df.Muon.pt)))
mc_kspread = rochester.kSpreadMC(
df.Muon.charge[hasgen],
df.Muon.pt[hasgen],
df.Muon.eta[hasgen],
df.Muon.phi[hasgen],
df.Muon.matched_gen.pt[hasgen],
)
mc_ksmear = rochester.kSmearMC(
df.Muon.charge[~hasgen],
df.Muon.pt[~hasgen],
df.Muon.eta[~hasgen],
df.Muon.phi[~hasgen],
df.Muon.nTrackerLayers[~hasgen],
mc_rand[~hasgen],
)
errspread = rochester.kSpreadMCerror(
df.Muon.charge[hasgen],
df.Muon.pt[hasgen],
df.Muon.eta[hasgen],
df.Muon.phi[hasgen],
df.Muon.matched_gen.pt[hasgen],
)
errsmear = rochester.kSmearMCerror(
df.Muon.charge[~hasgen],
df.Muon.pt[~hasgen],
df.Muon.eta[~hasgen],
df.Muon.phi[~hasgen],
df.Muon.nTrackerLayers[~hasgen],
mc_rand[~hasgen],
)
hasgen_flat = np.array(ak.flatten(hasgen))
corrections[hasgen_flat] = np.array(ak.flatten(mc_kspread))
corrections[~hasgen_flat] = np.array(ak.flatten(mc_ksmear))
errors[hasgen_flat] = np.array(ak.flatten(errspread))
errors[~hasgen_flat] = np.array(ak.flatten(errsmear))
corrections = ak.unflatten(corrections, ak.num(df.Muon.pt, axis=1))
errors = ak.unflatten(errors, ak.num(df.Muon.pt, axis=1))
else:
corrections = rochester.kScaleDT(df.Muon.charge, df.Muon.pt,
df.Muon.eta, df.Muon.phi)
errors = rochester.kScaleDTerror(df.Muon.charge, df.Muon.pt,
df.Muon.eta, df.Muon.phi)
df["Muon", "pt_roch"] = df.Muon.pt * corrections
df["Muon", "pt_roch_up"] = df.Muon.pt_roch + df.Muon.pt * errors
df["Muon", "pt_roch_down"] = df.Muon.pt_roch - df.Muon.pt * errors
def MT(leptons, met, debug=False):
if debug: set_trace()
# broadcast met into same shape as leptons
met_pt = (ak.ones_like(leptons.pt)) * (met.pt)
met_px = (ak.ones_like(leptons.pt)) * (met.px)
met_py = (ak.ones_like(leptons.pt)) * (met.py)
return np.sqrt(
np.square(leptons.pt + met_pt) - np.square(leptons.px + met_px) -
np.square(leptons.py + met_py))
def _kExtra(self, kpt, eta, nl, u, s=0, m=0):
# if it is a jagged array, save the offsets then flatten everything
# needed for the ternary conditions later
abseta = abs(eta)
kData = self._kRes[s][m][1](abseta) # type 1 is data
kMC = self._kRes[s][m][0](abseta) # type 0 is MC
mask = kData > kMC
x = awkward.zeros_like(kpt)
sigma = self._sigma(kpt, eta, nl, s, m)
# Rochester cbA = beta, cbN = m, as well as cbM (always 0?) = loc and cbS = scale to transform y = (x-loc)/scale in the pdf method
cbA = self._cbA[s][m](abseta, nl)
cbN = self._cbN[s][m](abseta, nl)
cbS = self._cbS[s][m](abseta, nl)
counts = awkward.num(u)
u_flat = awkward.flatten(u)
loc = awkward.zeros_like(u_flat)
cbA_flat = awkward.flatten(cbA)
cbN_flat = awkward.flatten(cbN)
cbS_flat = awkward.flatten(cbS)
invcdf = awkward.unflatten(
doublecrystalball.ppf(u_flat, cbA_flat, cbA_flat, cbN_flat,
cbN_flat, loc, cbS_flat),
counts,
)
x = awkward.where(
mask,
(numpy.sqrt(kData * kData - kMC * kMC) * sigma * invcdf),
x,
)
result = awkward.where(x > -1, 1.0 / (1.0 + x), awkward.ones_like(kpt))
if isinstance(kpt, numpy.ndarray):
result = numpy.array(result)
return result
def ak_solve_4PJ(self, mthad, mwhad, nschi):
"Inputs: m(thad_proxy), m(whad_proxy), and the NS solver sqrt(chi2) values (mthad, mwhad, nschi)\nOutputs: numpy array of the mass and NS disriminants (MassDiscr, NuDiscr)"
#set_trace()
mass_discval = ak.where((mthad < self.WTmass_right._axes[0][-1]) &
(mwhad < self.WTmass_right._axes[1][-1]),
self.WTmass_right(mthad, mwhad) /
np.sum(self.WTmass_right._values),
ak.ones_like(mthad) * np.nan)
masstest = -1 * np.log(mass_discval)
masstest = ak.nan_to_num(masstest, nan=np.inf)
nu_discval = ak.where(
nschi < self.NS_4PJ_right._axes[-1],
self.NS_4PJ_right(nschi) / np.sum(self.NS_4PJ_right._values),
ak.ones_like(nschi) * np.nan)
nstest = -1 * np.log(nu_discval)
nstest = ak.nan_to_num(nstest, nan=np.inf)
return masstest, nstest
def _assignTruthDef(self, tree):
assert self.splitIdx is not None
nonSplitRecHitSimClusIdx = self._createTruthAssociation(tree)
recHitTruthPID = self._assignTruthByIndexAndSplit(tree,"MergedSimCluster_pdgId",nonSplitRecHitSimClusIdx)
recHitTruthEnergy = self._assignTruthByIndexAndSplit(tree,"MergedSimCluster_boundaryEnergy",nonSplitRecHitSimClusIdx)
recHitDepEnergy = self._assignTruthByIndexAndSplit(tree,"MergedSimCluster_recEnergy",nonSplitRecHitSimClusIdx)
if not self.use_true_muon_momentum:
recHitTruthEnergy = ak1.where(np.abs(recHitTruthPID[:,:,0])==13, recHitDepEnergy, recHitTruthEnergy)
recHitTruthX = self._assignTruthByIndexAndSplit(tree,"MergedSimCluster_impactPoint_x",nonSplitRecHitSimClusIdx)
recHitTruthY = self._assignTruthByIndexAndSplit(tree,"MergedSimCluster_impactPoint_y",nonSplitRecHitSimClusIdx)
recHitTruthZ = self._assignTruthByIndexAndSplit(tree,"MergedSimCluster_impactPoint_z",nonSplitRecHitSimClusIdx)
recHitTruthTime = self._assignTruthByIndexAndSplit(tree,"MergedSimCluster_impactPoint_t",nonSplitRecHitSimClusIdx)
fullyContained = ak1.where(np.abs(recHitTruthZ)[:,:,0]<323.,#somehow that seems necessary
ak1.ones_like(recHitTruthZ),
ak1.zeros_like(recHitTruthZ))
recHitEnergy = self._readSplitAndExpand(tree,"RecHitHGC_energy")
recHitTime = self._readSplitAndExpand(tree,"RecHitHGC_time")
recHitX = self._readSplitAndExpand(tree,"RecHitHGC_x")
recHitY = self._readSplitAndExpand(tree,"RecHitHGC_y")
recHitZ = self._readSplitAndExpand(tree,"RecHitHGC_z")
# should not expand here to allow indexing as done below
recHitSimClusIdx = self._splitJaggedArray(nonSplitRecHitSimClusIdx)
# set noise to rec features
recHitTruthEnergy = ak1.where(recHitSimClusIdx<0, recHitEnergy, recHitTruthEnergy)
recHitTruthX = ak1.where(recHitSimClusIdx<0, recHitX, recHitTruthX)
recHitTruthY = ak1.where(recHitSimClusIdx<0, recHitY, recHitTruthY)
recHitTruthZ = ak1.where(recHitSimClusIdx<0, recHitZ, recHitTruthZ)
recHitTruthTime = ak1.where(recHitSimClusIdx<0, recHitTime, recHitTruthTime)
recHitDepEnergy = ak1.where(recHitSimClusIdx<0, recHitEnergy, recHitDepEnergy)
recHitSpectatorFlag = self._createSpectators(tree)
#remove spectator flag for noise
recHitSpectatorFlag = ak1.where(recHitSimClusIdx<0 , ak1.zeros_like(recHitSpectatorFlag), recHitSpectatorFlag)#this doesn't work for some reason!
#DEBUG!!!
#ticlidx = self._readSplitAndExpand(tree,'RecHitHGC_TICLCandIdx')
self.truth={} #DEBUG!!!
self.truth['t_idx'] = self._expand(recHitSimClusIdx)# now expand to a trailing dimension
self.truth['t_energy'] = recHitTruthEnergy
self.truth['t_pos'] = ak1.concatenate([recHitTruthX, recHitTruthY,recHitTruthZ],axis=-1)
self.truth['t_time'] = recHitTruthTime
self.truth['t_pid'] = recHitTruthPID
self.truth['t_spectator'] = recHitSpectatorFlag
self.truth['t_fully_contained'] = fullyContained
self.truth['t_rec_energy'] = recHitDepEnergy
def __init__(self, output, ev, weight, selection=None):
'''
output: the accumulator object
ev: NanoEvent
weight: coffea analysis_tools Weights object
'''
self.ev = ev
self.weight = weight.weight()
self.output = output
self.selection = None
self.addRow('entry', (ak.ones_like(self.weight) == 1))
def ak_solve_3J_lost(self, m2jets, nschi):
"Inputs: m(j1+j2) and the NS solver sqrt(chi2) values (m2jets, nschi)\nOutputs: numpy array of the combined, mass, and NS disriminants (MassDiscr, NuDiscr)"
#set_trace()
mass_discval = ak.where(
m2jets < self.Mass_3J_Lost_right._axes[-1],
self.Mass_3J_Lost_right(m2jets) /
np.sum(self.Mass_3J_Lost_right._values),
ak.ones_like(m2jets) * np.nan)
masstest = -1 * np.log(mass_discval)
masstest = ak.nan_to_num(masstest, nan=np.inf)
nu_discval = ak.where(
nschi < self.NS_3J_Lost_right._axes[-1],
self.NS_3J_Lost_right(nschi) /
np.sum(self.NS_3J_Lost_right._values),
ak.ones_like(nschi) * np.nan)
nstest = -1 * np.log(nu_discval)
nstest = ak.nan_to_num(nstest, nan=np.inf)
return masstest, nstest
def _readTree(self, tree):
self._readSplits(tree, splitlabel='Track_HGCFront_z')
trackPt = self._readSplitAndExpand(tree,"Track_pt")
trackEta = self._readSplitAndExpand(tree,"Track_HGCFront_eta")
trackVertEta = self._readSplitAndExpand(tree,"Track_eta")
trackMom = trackPt * np.cosh(trackVertEta)
impactX = self._readSplitAndExpand(tree,"Track_HGCFront_x")
impactY = self._readSplitAndExpand(tree,"Track_HGCFront_y")
impactZ = self._readSplitAndExpand(tree,"Track_HGCFront_z")
chi2 = self._readSplitAndExpand(tree,"Track_normChiSq")
impactR = np.sqrt(impactX**2+impactY**2+impactZ**2)+1e-3
impactTheta = np.arccos(impactZ/impactR)
self.features = ak1.concatenate([
trackMom,
trackEta,
ak1.ones_like(trackMom), #indicator if it is track or not
impactTheta,
impactR,
impactX,
impactY,
impactZ,
ak1.zeros_like(trackMom),#no time info (yet,could be from MTD here)
chi2 #this is radius for hits, here chi2 for tracks, since it's kinda realted to the impact points resolution...
], axis=-1)
#this is just for bookkeeping, keep them the same as for hits
self.featurenames = [
'recHitEnergy',
'recHitEta',
'isTrack',
'recHitTheta',
'recHitR',
'recHitX',
'recHitY',
'recHitZ',
'recHitTime',
'recHitHitR'
]
def get_charge_parent(particle):
parent = find_first_parent(particle)
charge = ak.zeros_like(parent)
one = [[-11, -13, -15, -17, 24, 37], 1]
minus_one = [[11, 13, 15, 17, -24, -37], -1]
two_thirds = [[2, 4, 6, 8], 2 / 3]
minus_two_thirds = [[-2, -4, -6, -8], -2 / 3]
minus_one_third = [[1, 3, 5, 7], -1 / 3]
one_third = [[-1, -3, -5, -7], 1 / 3]
zero = [[12, 14, 16, 18, 9, 21, 22, 23, 25], 0]
charge_pairs = [
one, minus_one, two_thirds, minus_two_thirds, minus_one_third, zero
]
for pair in charge_pairs:
for ID in pair[0]:
charge = (parent == ID) * ak.ones_like(parent) * pair[1] + (
~(parent == ID)) * charge
return charge
def build(self, jets, lazy_cache):
if lazy_cache is None:
raise Exception(
"CorrectedJetsFactory requires a awkward-array cache to function correctly."
)
lazy_cache = awkward._util.MappingProxy.maybe_wrap(lazy_cache)
if not isinstance(jets, awkward.highlevel.Array):
raise Exception(
"'jets' must be an awkward > 1.0.0 array of some kind!")
fields = awkward.fields(jets)
if len(fields) == 0:
raise Exception(
"Empty record, please pass a jet object with at least {self.real_sig} defined!"
)
out = awkward.flatten(jets)
wrap = partial(awkward_rewrap,
like_what=jets,
gfunc=rewrap_recordarray)
scalar_form = awkward.without_parameters(
out[self.name_map["ptRaw"]]).layout.form
in_dict = {field: out[field] for field in fields}
out_dict = dict(in_dict)
# take care of nominal JEC (no JER if available)
out_dict[self.name_map["JetPt"] +
"_orig"] = out_dict[self.name_map["JetPt"]]
out_dict[self.name_map["JetMass"] +
"_orig"] = out_dict[self.name_map["JetMass"]]
if self.treat_pt_as_raw:
out_dict[self.name_map["ptRaw"]] = out_dict[self.name_map["JetPt"]]
out_dict[self.name_map["massRaw"]] = out_dict[
self.name_map["JetMass"]]
jec_name_map = dict(self.name_map)
jec_name_map["JetPt"] = jec_name_map["ptRaw"]
jec_name_map["JetMass"] = jec_name_map["massRaw"]
if self.jec_stack.jec is not None:
jec_args = {
k: out_dict[jec_name_map[k]]
for k in self.jec_stack.jec.signature
}
out_dict[
"jet_energy_correction"] = self.jec_stack.jec.getCorrection(
**jec_args, form=scalar_form, lazy_cache=lazy_cache)
else:
out_dict["jet_energy_correction"] = awkward.without_parameters(
awkward.ones_like(out_dict[self.name_map["JetPt"]]))
# finally the lazy binding to the JEC
init_pt = partial(
awkward.virtual,
operator.mul,
args=(out_dict["jet_energy_correction"],
out_dict[self.name_map["ptRaw"]]),
cache=lazy_cache,
)
init_mass = partial(
awkward.virtual,
operator.mul,
args=(
out_dict["jet_energy_correction"],
out_dict[self.name_map["massRaw"]],
),
cache=lazy_cache,
)
out_dict[self.name_map["JetPt"]] = init_pt(length=len(out),
form=scalar_form)
out_dict[self.name_map["JetMass"]] = init_mass(length=len(out),
form=scalar_form)
out_dict[self.name_map["JetPt"] +
"_jec"] = out_dict[self.name_map["JetPt"]]
out_dict[self.name_map["JetMass"] +
"_jec"] = out_dict[self.name_map["JetMass"]]
# in jer we need to have a stash for the intermediate JEC products
has_jer = False
if self.jec_stack.jer is not None and self.jec_stack.jersf is not None:
has_jer = True
jer_name_map = dict(self.name_map)
jer_name_map["JetPt"] = jer_name_map["JetPt"] + "_jec"
jer_name_map["JetMass"] = jer_name_map["JetMass"] + "_jec"
jerargs = {
k: out_dict[jer_name_map[k]]
for k in self.jec_stack.jer.signature
}
out_dict[
"jet_energy_resolution"] = self.jec_stack.jer.getResolution(
**jerargs, form=scalar_form, lazy_cache=lazy_cache)
jersfargs = {
k: out_dict[jer_name_map[k]]
for k in self.jec_stack.jersf.signature
}
out_dict[
"jet_energy_resolution_scale_factor"] = self.jec_stack.jersf.getScaleFactor(
**jersfargs, form=_JERSF_FORM, lazy_cache=lazy_cache)
seeds = numpy.array(out_dict[self.name_map["JetPt"] +
"_orig"])[[0, -1]].view("i4")
out_dict["jet_resolution_rand_gauss"] = awkward.virtual(
rand_gauss,
args=(
out_dict[self.name_map["JetPt"] + "_orig"],
numpy.random.Generator(numpy.random.PCG64(seeds)),
),
cache=lazy_cache,
length=len(out),
form=scalar_form,
)
init_jerc = partial(
awkward.virtual,
jer_smear,
args=(
0,
self.forceStochastic,
out_dict[jer_name_map["ptGenJet"]],
out_dict[jer_name_map["JetPt"]],
out_dict[jer_name_map["JetEta"]],
out_dict["jet_energy_resolution"],
out_dict["jet_resolution_rand_gauss"],
out_dict["jet_energy_resolution_scale_factor"],
),
cache=lazy_cache,
)
out_dict["jet_energy_resolution_correction"] = init_jerc(
length=len(out), form=scalar_form)
init_pt_jer = partial(
awkward.virtual,
operator.mul,
args=(
out_dict["jet_energy_resolution_correction"],
out_dict[jer_name_map["JetPt"]],
),
cache=lazy_cache,
)
init_mass_jer = partial(
awkward.virtual,
operator.mul,
args=(
out_dict["jet_energy_resolution_correction"],
out_dict[jer_name_map["JetMass"]],
),
cache=lazy_cache,
)
out_dict[self.name_map["JetPt"]] = init_pt_jer(length=len(out),
form=scalar_form)
out_dict[self.name_map["JetMass"]] = init_mass_jer(
length=len(out), form=scalar_form)
out_dict[self.name_map["JetPt"] +
"_jer"] = out_dict[self.name_map["JetPt"]]
out_dict[self.name_map["JetMass"] +
"_jer"] = out_dict[self.name_map["JetMass"]]
# JER systematics
jerc_up = partial(
awkward.virtual,
jer_smear,
args=(
1,
self.forceStochastic,
out_dict[jer_name_map["ptGenJet"]],
out_dict[jer_name_map["JetPt"]],
out_dict[jer_name_map["JetEta"]],
out_dict["jet_energy_resolution"],
out_dict["jet_resolution_rand_gauss"],
out_dict["jet_energy_resolution_scale_factor"],
),
cache=lazy_cache,
)
up = awkward.flatten(jets)
up["jet_energy_resolution_correction"] = jerc_up(length=len(out),
form=scalar_form)
init_pt_jer = partial(
awkward.virtual,
operator.mul,
args=(
up["jet_energy_resolution_correction"],
out_dict[jer_name_map["JetPt"]],
),
cache=lazy_cache,
)
init_mass_jer = partial(
awkward.virtual,
operator.mul,
args=(
up["jet_energy_resolution_correction"],
out_dict[jer_name_map["JetMass"]],
),
cache=lazy_cache,
)
up[self.name_map["JetPt"]] = init_pt_jer(length=len(out),
form=scalar_form)
up[self.name_map["JetMass"]] = init_mass_jer(length=len(out),
form=scalar_form)
jerc_down = partial(
awkward.virtual,
jer_smear,
args=(
2,
self.forceStochastic,
out_dict[jer_name_map["ptGenJet"]],
out_dict[jer_name_map["JetPt"]],
out_dict[jer_name_map["JetEta"]],
out_dict["jet_energy_resolution"],
out_dict["jet_resolution_rand_gauss"],
out_dict["jet_energy_resolution_scale_factor"],
),
cache=lazy_cache,
)
down = awkward.flatten(jets)
down["jet_energy_resolution_correction"] = jerc_down(
length=len(out), form=scalar_form)
init_pt_jer = partial(
awkward.virtual,
operator.mul,
args=(
down["jet_energy_resolution_correction"],
out_dict[jer_name_map["JetPt"]],
),
cache=lazy_cache,
)
init_mass_jer = partial(
awkward.virtual,
operator.mul,
args=(
down["jet_energy_resolution_correction"],
out_dict[jer_name_map["JetMass"]],
),
cache=lazy_cache,
)
down[self.name_map["JetPt"]] = init_pt_jer(length=len(out),
form=scalar_form)
down[self.name_map["JetMass"]] = init_mass_jer(length=len(out),
form=scalar_form)
out_dict["JER"] = awkward.zip({
"up": up,
"down": down
},
depth_limit=1,
with_name="JetSystematic")
if self.jec_stack.junc is not None:
juncnames = {}
juncnames.update(self.name_map)
if has_jer:
juncnames["JetPt"] = juncnames["JetPt"] + "_jer"
juncnames["JetMass"] = juncnames["JetMass"] + "_jer"
else:
juncnames["JetPt"] = juncnames["JetPt"] + "_jec"
juncnames["JetMass"] = juncnames["JetMass"] + "_jec"
juncargs = {
k: out_dict[juncnames[k]]
for k in self.jec_stack.junc.signature
}
juncs = self.jec_stack.junc.getUncertainty(**juncargs)
def junc_smeared_val(uncvals, up_down, variable):
return awkward.materialized(uncvals[:, up_down] * variable)
def build_variation(unc, jetpt, jetpt_orig, jetmass, jetmass_orig,
updown):
var_dict = dict(in_dict)
var_dict[jetpt] = awkward.virtual(
junc_smeared_val,
args=(
unc,
updown,
jetpt_orig,
),
length=len(out),
form=scalar_form,
cache=lazy_cache,
)
var_dict[jetmass] = awkward.virtual(
junc_smeared_val,
args=(
unc,
updown,
jetmass_orig,
),
length=len(out),
form=scalar_form,
cache=lazy_cache,
)
return awkward.zip(
var_dict,
depth_limit=1,
parameters=out.layout.parameters,
behavior=out.behavior,
)
def build_variant(unc, jetpt, jetpt_orig, jetmass, jetmass_orig):
up = build_variation(unc, jetpt, jetpt_orig, jetmass,
jetmass_orig, 0)
down = build_variation(unc, jetpt, jetpt_orig, jetmass,
jetmass_orig, 1)
return awkward.zip({
"up": up,
"down": down
},
depth_limit=1,
with_name="JetSystematic")
for name, func in juncs:
out_dict[f"jet_energy_uncertainty_{name}"] = func
out_dict[f"JES_{name}"] = build_variant(
func,
self.name_map["JetPt"],
out_dict[juncnames["JetPt"]],
self.name_map["JetMass"],
out_dict[juncnames["JetMass"]],
)
out_parms = out.layout.parameters
out_parms["corrected"] = True
out = awkward.zip(out_dict,
depth_limit=1,
parameters=out_parms,
behavior=out.behavior)
return wrap(out)
def test():
array = ak.Array(
[
[{"x": 0.0, "y": []}, {"x": 1.1, "y": [1]}, {"x": 2.2, "y": [1, 2]}],
[],
[
{"x": 3.3, "y": [1, 2, None, 3]},
False,
False,
True,
{"x": 4.4, "y": [1, 2, None, 3, 4]},
],
]
)
def assert_array_type(new_array, intended_type):
"""helper function to check each part of the array took the intended type"""
assert isinstance(new_array[0][0]["x"], intended_type)
assert isinstance(new_array[0][1]["x"], intended_type)
assert isinstance(new_array[0][1]["y"][0], intended_type)
assert isinstance(new_array[0][2]["x"], intended_type)
assert isinstance(new_array[0][2]["y"][0], intended_type)
assert isinstance(new_array[0][2]["y"][1], intended_type)
assert isinstance(new_array[2][0]["x"], intended_type)
assert isinstance(new_array[2][0]["y"][0], intended_type)
assert isinstance(new_array[2][0]["y"][1], intended_type)
assert isinstance(new_array[2][0]["y"][3], intended_type)
assert isinstance(new_array[2][1], intended_type)
assert isinstance(new_array[2][2], intended_type)
assert isinstance(new_array[2][3], intended_type)
assert isinstance(new_array[2][4]["x"], intended_type)
assert isinstance(new_array[2][4]["y"][0], intended_type)
assert isinstance(new_array[2][4]["y"][1], intended_type)
assert isinstance(new_array[2][4]["y"][3], intended_type)
assert isinstance(new_array[2][4]["y"][4], intended_type)
int_type = ak.full_like(array, 12, dtype=int)
float_type = ak.full_like(array, 12, dtype=float)
assert int_type.tolist() == float_type.tolist()
assert_array_type(int_type, int)
assert_array_type(float_type, float)
bool_type = ak.full_like(array, 12, dtype=bool)
assert_array_type(bool_type, bool)
int_type = ak.full_like(array, -1.2, dtype=int)
float_type = ak.full_like(array, -1.2, dtype=float)
bool_type = ak.full_like(array, -1.2, dtype=bool)
assert_array_type(int_type, int)
assert_array_type(float_type, float)
assert_array_type(bool_type, bool)
int_type = ak.zeros_like(array, dtype=int)
float_type = ak.zeros_like(array, dtype=float)
bool_type = ak.zeros_like(array, dtype=bool)
assert int_type.tolist() == float_type.tolist()
assert int_type.tolist() == bool_type.tolist()
assert_array_type(int_type, int)
assert_array_type(float_type, float)
assert_array_type(bool_type, bool)
int_type = ak.ones_like(array, dtype=int)
float_type = ak.ones_like(array, dtype=float)
bool_type = ak.ones_like(array, dtype=bool)
assert int_type.tolist() == float_type.tolist()
assert int_type.tolist() == bool_type.tolist()
assert_array_type(int_type, int)
assert_array_type(float_type, float)
assert_array_type(bool_type, bool)
array = ak.Array([["one", "two", "three"], [], ["four", "five"]])
def assert_array_type(new_array, intended_type):
"""helper function to check each part of the array took the intended type"""
assert isinstance(new_array[0][0], intended_type)
assert isinstance(new_array[0][1], intended_type)
assert isinstance(new_array[0][2], intended_type)
assert isinstance(new_array[2][0], intended_type)
assert isinstance(new_array[2][1], intended_type)
int_type = ak.full_like(array, 12, dtype=int)
float_type = ak.full_like(array, 12, dtype=float)
assert int_type.tolist() == float_type.tolist()
assert_array_type(int_type, int)
assert_array_type(float_type, float)
bool_type = ak.full_like(array, 12, dtype=bool)
assert_array_type(bool_type, bool)
int_type = ak.full_like(array, -1.2, dtype=int)
float_type = ak.full_like(array, -1.2, dtype=float)
bool_type = ak.full_like(array, -1.2, dtype=bool)
assert_array_type(int_type, int)
assert_array_type(float_type, float)
assert_array_type(bool_type, bool)
int_type = ak.zeros_like(array, dtype=int)
float_type = ak.zeros_like(array, dtype=float)
bool_type = ak.zeros_like(array, dtype=bool)
assert int_type.tolist() == float_type.tolist()
assert int_type.tolist() == bool_type.tolist()
assert_array_type(int_type, int)
assert_array_type(float_type, float)
assert_array_type(bool_type, bool)
int_type = ak.ones_like(array, dtype=int)
float_type = ak.ones_like(array, dtype=float)
bool_type = ak.ones_like(array, dtype=bool)
assert int_type.tolist() == float_type.tolist()
assert int_type.tolist() == bool_type.tolist()
assert_array_type(int_type, int)
assert_array_type(float_type, float)
assert_array_type(bool_type, bool)
def test():
array = ak.Array([
[{
"x": 0.0,
"y": []
}, {
"x": 1.1,
"y": [1]
}, {
"x": 2.2,
"y": [1, 2]
}],
[],
[
{
"x": 3.3,
"y": [1, 2, None, 3]
},
False,
False,
True,
{
"x": 4.4,
"y": [1, 2, None, 3, 4]
},
],
])
assert ak.full_like(array, 12.3).tolist() == [
[{
"x": 12.3,
"y": []
}, {
"x": 12.3,
"y": [12]
}, {
"x": 12.3,
"y": [12, 12]
}],
[],
[
{
"x": 12.3,
"y": [12, 12, None, 12]
},
True,
True,
True,
{
"x": 12.3,
"y": [12, 12, None, 12, 12]
},
],
]
assert ak.zeros_like(array).tolist() == [
[{
"x": 0.0,
"y": []
}, {
"x": 0.0,
"y": [0]
}, {
"x": 0.0,
"y": [0, 0]
}],
[],
[
{
"x": 0.0,
"y": [0, 0, None, 0]
},
False,
False,
False,
{
"x": 0.0,
"y": [0, 0, None, 0, 0]
},
],
]
assert ak.ones_like(array).tolist() == [
[{
"x": 1.0,
"y": []
}, {
"x": 1.0,
"y": [1]
}, {
"x": 1.0,
"y": [1, 1]
}],
[],
[
{
"x": 1.0,
"y": [1, 1, None, 1]
},
True,
True,
True,
{
"x": 1.0,
"y": [1, 1, None, 1, 1]
},
],
]
array = ak.Array([["one", "two", "three"], [], ["four", "five"]])
assert ak.full_like(array, "hello").tolist() == [
["hello", "hello", "hello"],
[],
["hello", "hello"],
]
assert ak.full_like(array, 1).tolist() == [["1", "1", "1"], [], ["1", "1"]]
assert ak.full_like(array, 0).tolist() == [["0", "0", "0"], [], ["0", "0"]]
assert ak.ones_like(array).tolist() == [["1", "1", "1"], [], ["1", "1"]]
assert ak.zeros_like(array).tolist() == [["", "", ""], [], ["", ""]]
array = ak.Array([[b"one", b"two", b"three"], [], [b"four", b"five"]])
assert ak.full_like(array, b"hello").tolist() == [
[b"hello", b"hello", b"hello"],
[],
[b"hello", b"hello"],
]
assert ak.full_like(array, 1).tolist() == [[b"1", b"1", b"1"], [],
[b"1", b"1"]]
assert ak.full_like(array, 0).tolist() == [[b"0", b"0", b"0"], [],
[b"0", b"0"]]
assert ak.ones_like(array).tolist() == [[b"1", b"1", b"1"], [],
[b"1", b"1"]]
assert ak.zeros_like(array).tolist() == [[b"", b"", b""], [], [b"", b""]]
def __init__(self, ev, obj, wp, year=2018, verbose=0):
self.obj = obj
self.wp = wp
if self.wp == None:
self.selection_dict = {}
else:
self.selection_dict = obj_def[self.obj][self.wp]
self.v = verbose
self.year = year
id_level = None
if wp.lower().count('veto') or wp.lower().count('loose'):
id_level = 0
elif wp.lower().count('fake'):
id_level = 1
elif wp.lower().count('tight'):
id_level = 2
if self.obj == "Muon":
# collections are already there, so we just need to calculate missing ones
ev['Muon', 'absMiniIso'] = ev.Muon.miniPFRelIso_all * ev.Muon.pt
ev['Muon', 'ptErrRel'] = ev.Muon.ptErr / ev.Muon.pt
# this is what we are using:
# - jetRelIso if the matched jet is within deltaR<0.4, pfRelIso03_all otherwise
# - btagDeepFlavB discriminator of the matched jet if jet is within deltaR<0.4, 0 otherwise
# - pt_cone = 0.9*pt of matched jet if jet is within deltaR<0.4, pt/(pt+iso) otherwise
mask_close = (ak.fill_none(ev.Muon.delta_r(ev.Muon.matched_jet),
99) < 0.4) * 1
mask_far = ~(ak.fill_none(ev.Muon.delta_r(ev.Muon.matched_jet), 99)
< 0.4) * 1
deepJet = ak.fill_none(ev.Muon.matched_jet.btagDeepFlavB,
0) * mask_close + 0 * mask_far
jetRelIsoV2 = ev.Muon.jetRelIso * mask_close + ev.Muon.pfRelIso03_all * mask_far # default to 0 if no match
#conePt = 0.9 * ak.fill_none(ev.Muon.matched_jet.pt,0) * mask_close + ev.Muon.pt*(1 + ev.Muon.miniPFRelIso_all)*mask_far
if self.year == 2017 or self.year == 2018:
I_1 = 0.11
I_2 = 0.74
I_3 = 6.8
elif self.year == 2016:
I_1 = 0.16
I_2 = 0.76
I_3 = 7.2
PF_unflatten = ak.from_regular(
ev.Muon.miniPFRelIso_all[:, :, np.newaxis])
max_miniIso = ak.max(
ak.concatenate(
[PF_unflatten - I_1,
ak.zeros_like(PF_unflatten)], axis=2),
axis=2) #equivalent to max(0, ev.Muon.miniPFRelIso_all - I_1)
muon_pt_unflatten = ak.from_regular(ev.Muon.pt[:, :, np.newaxis])
jet_pt_unflatten = ak.from_regular(
ev.Muon.matched_jet.pt[:, :, np.newaxis])
max_pt = ak.max(
ak.concatenate([muon_pt_unflatten, jet_pt_unflatten * I_2],
axis=2),
axis=2) #max(ev.Muon.pt, ev.Muon.matched_jet.pt * I_2)
conePt = (ev.Muon.pt *
(1 + max_miniIso)) * (ev.Muon.jetPtRelv2 > I_3) + (
max_pt * ~(ev.Muon.jetPtRelv2 > I_3))
ev['Muon', 'deepJet'] = ak.copy(deepJet)
ev['Muon', 'jetRelIsoV2'] = jetRelIsoV2
ev['Muon', 'conePt'] = conePt
ev['Muon', 'id'] = ak.ones_like(conePt) * id_level
self.cand = ev.Muon
elif self.obj == "Electron":
# calculate new variables. asignment is awkward, but what can you do.
ev['Electron',
'absMiniIso'] = ev.Electron.miniPFRelIso_all * ev.Electron.pt
ev['Electron', 'etaSC'] = ev.Electron.eta + ev.Electron.deltaEtaSC
if self.year == 2017 or self.year == 2018:
I_1 = 0.07
I_2 = 0.78
I_3 = 8.0
elif self.year == 2016:
I_1 = 0.12
I_2 = 0.80
I_3 = 7.2
# the following line is only needed if we do our own matching.
# right now, we keep using the NanoAOD match, but check the deltaR distance
# jet_index, mask_match, mask_nomatch = self.matchJets(ev.Electron, ev.Jet)
# this is what we are using:
# - jetRelIso if the matched jet is within deltaR<0.4, pfRelIso03_all otherwise
# - btagDeepFlavB discriminator of the matched jet if jet is within deltaR<0.4, 0 otherwise
# - pt_cone = 0.9*pt of matched jet if jet is within deltaR<0.4, pt/(pt+iso) otherwise
mask_close = (ak.fill_none(
ev.Electron.delta_r(ev.Electron.matched_jet), 99) < 0.4) * 1
mask_far = ~(ak.fill_none(
ev.Electron.delta_r(ev.Electron.matched_jet), 99) < 0.4) * 1
deepJet = ak.fill_none(ev.Electron.matched_jet.btagDeepFlavB,
0) * mask_close
jetRelIsoV2 = ev.Electron.jetRelIso * mask_close + ev.Electron.pfRelIso03_all * mask_far # default to 0 if no match
#conePt = 0.9 * ak.fill_none(ev.Electron.matched_jet.pt,0) * mask_close + ev.Electron.pt*(1 + ev.Electron.miniPFRelIso_all)*mask_far
PF_unflatten = ak.from_regular(
ev.Electron.miniPFRelIso_all[:, :, np.newaxis])
max_miniIso = ak.max(
ak.concatenate(
[PF_unflatten - I_1,
ak.zeros_like(PF_unflatten)], axis=2),
axis=2) #equivalent to max(0, ev.Muon.miniPFRelIso_all - I_1)
electron_pt_unflatten = ak.from_regular(ev.Electron.pt[:, :,
np.newaxis])
jet_pt_unflatten = ak.from_regular(
ev.Electron.matched_jet.pt[:, :, np.newaxis])
max_pt = ak.max(
ak.concatenate([electron_pt_unflatten, jet_pt_unflatten * I_2],
axis=2),
axis=2) #max(ev.Muon.pt, ev.Muon.matched_jet.pt * I_2)
conePt = (ev.Electron.pt *
(1 + max_miniIso)) * (ev.Electron.jetPtRelv2 > I_3) + (
max_pt * ~(ev.Electron.jetPtRelv2 > I_3))
ev['Electron', 'deepJet'] = ak.copy(deepJet)
ev['Electron', 'jetRelIsoV2'] = jetRelIsoV2
ev['Electron', 'conePt'] = conePt
ev['Electron', 'id'] = ak.ones_like(conePt) * id_level
ev['Electron', 'jetRelIso'] = ev.Electron.jetRelIso
ev['Electron', 'jetPtRelv2'] = ev.Electron.jetPtRelv2
self.cand = ev.Electron
self.getSelection()
if self.obj == "Electron" and self.wp == "tight":
self.selection = self.selection & self.getElectronMVAID(
) & self.getIsolation(0.07, 0.78, 8.0) & self.isTriggerSafeNoIso()
if self.v > 0: print(" - custom ID and multi-isolation")
if self.obj == "Electron" and self.wp == "tightFCNC":
self.selection = self.selection & self.getElectronMVAID(
) & self.getFCNCIsolation(ev.Electron.jetRelIso,
ev.Electron.jetPtRelv2, I_2,
I_3) & (ev.Electron.miniPFRelIso_all <
I_1) & self.isTriggerSafeNoIso()
if self.v > 0: print(" - custom ID and multi-isolation")
if self.obj == "Muon" and self.wp == "tight":
self.selection = self.selection & self.getIsolation(
0.11, 0.74, 6.8)
if self.v > 0: print(" - custom multi-isolation")
#self.selection = self.selection & ak.fill_none(ev.Muon.matched_jet.btagDeepFlavB<0.2770, True)
#self.selection = self.selection & (ev.Muon.matched_jet.btagDeepFlavB<0.2770)
#if self.v>0: print (" - deepJet")
if self.obj == "Muon" and self.wp == "tightFCNC":
self.selection = self.selection & self.getFCNCIsolation(
ev.Muon.jetRelIso, ev.Muon.jetPtRelv2, I_2,
I_3) & (ev.Muon.miniPFRelIso_all < I_1)
if self.v > 0: print(" - custom multi-isolation")
if self.obj == "Electron" and (self.wp == "tightTTH"
or self.wp == 'fakeableTTH'
or self.wp == "tightSSTTH"
or self.wp == 'fakeableSSTTH'):
self.selection = self.selection & self.getSigmaIEtaIEta
if self.v > 0: print(" - SigmaIEtaIEta")
#self.selection = self.selection & ak.fill_none(ev.Electron.matched_jet.btagDeepFlavB<0.2770, True)
#self.selection = self.selection & (ev.Electron.matched_jet.btagDeepFlavB<0.2770)
#self.selection = self.selection & (ev.Jet[ev.Electron.jetIdx].btagDeepFlavB<0.2770)
#if self.v>0: print (" - deepJet")
if self.obj == "Electron" and self.wp == "looseFCNC":
self.selection = self.selection & (ev.Electron.miniPFRelIso_all <
0.4)
if self.obj == 'Muon' and (self.wp == 'fakeableTTH'
or self.wp == 'fakeableSSTTH'):
#self.selection = self.selection & (self.cand.deepJet < self.getThreshold(self.cand.conePt, min_pt=20, max_pt=45, low=0.2770, high=0.0494))
self.selection = self.selection & (ak.fill_none(
ev.Muon.matched_jet.btagDeepFlavB, 0) < self.getThreshold(
self.cand.conePt, min_pt=20, max_pt=45))
if self.v > 0: print(" - interpolated deepJet")
if self.obj == "Muon" and self.wp == "looseFCNC":
self.selection = self.selection & (ev.Muon.miniPFRelIso_all < 0.4)
def _assignTruthCPPU(self, tree):
'''
'''
print('\n\n>>>>>>WARNING: using calo particle plus PU mode: unless this is a special testing sample, you should not be using this function!<<<<\n\n')
print('\n\n>>>>>> The configuration here also assumes exactly 1 non PU particle per endcap!<<<<<<<\n\n')
assert self.splitIdx is not None
scidx = self._readArray(tree,"RecHitHGC_BestSimClusterIdx")
cpidx = self._readArray(tree,"SimCluster_CaloPartIdx")
nonSplitRecHitSimClusIdx = ak1.where(scidx >= 0, cpidx[scidx], -1)#mark noise
non_pu_cp = (self._readArray(tree,"CaloPart_eventId")+self._readArray(tree,"CaloPart_bunchCrossing"))<1
recHitNotPU = self._assignTruthByIndexAndSplit(tree,non_pu_cp,nonSplitRecHitSimClusIdx)
recHitTruthPID = self._assignTruthByIndexAndSplit(tree,"CaloPart_pdgId",nonSplitRecHitSimClusIdx)
recHitTruthEnergy = self._assignTruthByIndexAndSplit(tree,"CaloPart_energy",nonSplitRecHitSimClusIdx)
fzeros = self._assignTruthByIndexAndSplit(tree,"CaloPart_pt",nonSplitRecHitSimClusIdx)*0.
recHitTruthX = fzeros
recHitTruthY = fzeros
recHitTruthZ = fzeros
recHitTruthTime = fzeros
recHitDepEnergy = fzeros
fullyContained = ak1.ones_like(fzeros)
from globals import pu
recHitSimClusIdx = self._splitJaggedArray(nonSplitRecHitSimClusIdx)
recHitSimClusIdx = self._expand(recHitSimClusIdx)
recHitSimClusIdx = recHitSimClusIdx + pu.t_idx_offset*(1-recHitNotPU)*(recHitSimClusIdx>=0)#only for not noise
#pu_idx_offset[recHitNotPU[...,0]>0]=0
#recHitSimClusIdx += pu_idx_offset
#recHitSimClusIdx = recHitSimClusIdx + pu.t_idx_offset * (1-recHitNotPU[...,0])
#testing
#recHitSimClusIdx = recHitNotPU #(1+recHitSimClusIdx)*recHitNotPU[...,0]#[...,0]
recHitSpectatorFlag = fzeros
self.truth={}
self.truth['t_idx'] = recHitSimClusIdx
self.truth['t_energy'] = recHitTruthEnergy
self.truth['t_pos'] = ak1.concatenate([recHitTruthX, recHitTruthY,recHitTruthZ],axis=-1)
self.truth['t_time'] = recHitTruthTime
self.truth['t_pid'] = recHitTruthPID
self.truth['t_spectator'] = recHitSpectatorFlag
self.truth['t_fully_contained'] = fullyContained
self.truth['t_rec_energy'] = recHitDepEnergy
if not self.cp_plus_pu_mode_reduce:
return
print('\n>>>reducing set by ∆R: handle with care<<<\n')
#return
#now remove access particles around initial one
#this might not work because of close by gun
#.... it does not work....
#npu_cpeta = self._readArray(tree,"CaloPart_eta")
#npu_cpphi = self._readArray(tree,"CaloPart_phi")
recHitX = self._readSplitAndExpand(tree,"RecHitHGC_x")# EC x V x 1
recHitY = self._readSplitAndExpand(tree,"RecHitHGC_y")
recHitZ = self._readSplitAndExpand(tree,"RecHitHGC_z")
cp_x = ak1.mean(recHitX[recHitNotPU[...,0]>0],axis=1)
cp_y = ak1.mean(recHitY[recHitNotPU[...,0]>0],axis=1)
cp_z = ak1.mean(recHitZ[recHitNotPU[...,0]>0],axis=1)
recHitEta = calc_eta(recHitX, recHitY, recHitZ)
recHitPhi = calc_phi(recHitX, recHitY, recHitZ)
#return
# ...
#just loop over EC, there are only very few
hitselector=[]
for i_endcap in range(len(recHitEta)):
ecps_eta=calc_eta(cp_x[i_endcap], cp_y[i_endcap], cp_z[i_endcap])
ecps_phi=calc_phi(cp_x[i_endcap], cp_y[i_endcap], cp_z[i_endcap])
ec_heta = recHitEta[i_endcap].to_numpy() # V x 1
ec_hphi = recHitPhi[i_endcap].to_numpy()
deta = ec_heta-ecps_eta
dphi = deltaPhi(ec_hphi,ecps_phi) # V x 1
drsq = deta**2 + dphi**2 # V x 1
close = np.array(drsq < 0.5**2,dtype='int')[...,0]
#either_close = np.sum(close, axis=1)#does sum work on bool?
#print(close.shape)
hitselector.append(close)
#last
hitselector = ak1.from_iter(hitselector) # EC x V'
#hitselector = recHitNotPU[...,0]
#print('hitselector',ak1.sum(hitselector))
for k in self.truth.keys():
self.truth[k] = self.truth[k][hitselector>0]
self.features = self.features[hitselector>0]
def select_normal(genparts, w_decay_momid):
columns = ["pt", "eta", "phi", "mass", "pdgId", "charge", "decaytype"]
# get tops, defined as last copy
gen_tops = genparts[(genparts.hasFlags(['isLastCopy']))
& (genparts.pdgId == 6)]
gen_tops['charge'] = ak.ones_like(gen_tops.pt) * (2. / 3.)
gen_tbars = genparts[(genparts.hasFlags(['isLastCopy']))
& (genparts.pdgId == -6)]
gen_tbars['charge'] = ak.ones_like(gen_tbars.pt) * (-2. / 3.)
# get direct top decay products (children are "isHardProcess" and "isFirstCopy")
gen_bs = ak.flatten(gen_tops.children[(gen_tops.children.pdgId == 5)],
axis=2)
gen_bs['charge'] = ak.ones_like(gen_bs.pt) * (-1. / 3.)
gen_bbars = ak.flatten(
gen_tbars.children[(gen_tbars.children.pdgId == -5)], axis=2)
gen_bbars['charge'] = ak.ones_like(gen_bbars.pt) * (1. / 3.)
gen_wplus = ak.flatten(
gen_tops.children[(gen_tops.children.pdgId == w_decay_momid)], axis=2)
gen_wplus['charge'] = ak.ones_like(gen_wplus.pt)
gen_wminus = ak.flatten(
gen_tbars.children[(gen_tbars.children.pdgId == -1 * w_decay_momid)],
axis=2)
gen_wminus['charge'] = ak.ones_like(gen_wminus.pt) * (-1.)
# get w decay products
# get last copy of W bosons
last_gen_wplus = ak.flatten(gen_tops.distinctChildren[
(gen_tops.distinctChildren.pdgId == w_decay_momid)
& (gen_tops.distinctChildren.hasFlags(['isLastCopy']))],
axis=2)
last_gen_wminus = ak.flatten(gen_tbars.distinctChildren[
(gen_tbars.distinctChildren.pdgId == -1 * w_decay_momid)
& (gen_tbars.distinctChildren.hasFlags(['isLastCopy']))],
axis=2)
gen_wplus_decaytype = np.zeros(len(gen_wplus))
gen_wminus_decaytype = np.zeros(len(gen_wminus))
# up/down partons from last W+
gen_wpartons_up_fromWplus = ak.flatten(
last_gen_wplus.children[(np.mod(last_gen_wplus.children.pdgId, 2) == 0)
& (np.abs(last_gen_wplus.children.pdgId) < 6)],
axis=2)
gen_wpartons_up_fromWplus['charge'] = np.sign(
gen_wpartons_up_fromWplus.pdgId) * (2. / 3.)
gen_wplus_decaytype[ak.num(gen_wpartons_up_fromWplus) > 0] = np.ones(
ak.sum(ak.num(gen_wpartons_up_fromWplus) > 0)) * 2
gen_wpartons_dw_fromWplus = ak.flatten(
last_gen_wplus.children[(np.mod(last_gen_wplus.children.pdgId, 2) == 1)
& (np.abs(last_gen_wplus.children.pdgId) < 6)],
axis=2)
gen_wpartons_dw_fromWplus['charge'] = np.sign(
gen_wpartons_up_fromWplus.pdgId) * (-1. / 3.)
# up/down partons from last W-
gen_wpartons_up_fromWminus = ak.flatten(last_gen_wminus.children[
(np.mod(last_gen_wminus.children.pdgId, 2) == 0)
& (np.abs(last_gen_wminus.children.pdgId) < 6)],
axis=2)
gen_wpartons_up_fromWminus['charge'] = np.sign(
gen_wpartons_up_fromWminus.pdgId) * (2. / 3.)
gen_wminus_decaytype[ak.num(gen_wpartons_up_fromWminus) > 0] = np.ones(
ak.sum(ak.num(gen_wpartons_up_fromWminus) > 0)) * 2
gen_wpartons_dw_fromWminus = ak.flatten(last_gen_wminus.children[
(np.mod(last_gen_wminus.children.pdgId, 2) == 1)
& (np.abs(last_gen_wminus.children.pdgId) < 6)],
axis=2)
gen_wpartons_dw_fromWminus['charge'] = np.sign(
gen_wpartons_up_fromWminus.pdgId) * (-1. / 3.)
# charged leps from last W+
gen_charged_leps_fromWplus = ak.flatten(
last_gen_wplus.children[(np.abs(last_gen_wplus.children.pdgId) == 11) |
(np.abs(last_gen_wplus.children.pdgId) == 13) |
(np.abs(last_gen_wplus.children.pdgId) == 15)],
axis=2)
gen_charged_leps_fromWplus['charge'] = ak.ones_like(
gen_charged_leps_fromWplus.pdgId)
gen_wplus_decaytype[ak.num(gen_charged_leps_fromWplus) > 0] = np.ones(
ak.sum(ak.num(gen_charged_leps_fromWplus) > 0))
# add decaytype (0 is INVALID, 1 is LEPTONIC, 2 is HADRONIC)
gen_wplus['decaytype'] = ak.unflatten(gen_wplus_decaytype,
ak.num(gen_wplus))
gen_tops['decaytype'] = gen_wplus['decaytype'] # set decaytype for tops
gen_bs['decaytype'] = gen_wplus['decaytype'] # set decaytype for bs
# neutral leps from last W+
gen_neutral_leps_fromWplus = ak.flatten(
last_gen_wplus.children[(np.abs(last_gen_wplus.children.pdgId) == 12) |
(np.abs(last_gen_wplus.children.pdgId) == 14) |
(np.abs(last_gen_wplus.children.pdgId) == 16)],
axis=2)
gen_neutral_leps_fromWplus['charge'] = ak.zeros_like(
gen_neutral_leps_fromWplus.pt)
# charged leps from last W-
gen_charged_leps_fromWminus = ak.flatten(last_gen_wminus.children[
(np.abs(last_gen_wminus.children.pdgId) == 11) |
(np.abs(last_gen_wminus.children.pdgId) == 13) |
(np.abs(last_gen_wminus.children.pdgId) == 15)],
axis=2)
gen_charged_leps_fromWminus['charge'] = ak.ones_like(
gen_charged_leps_fromWminus.pdgId) * (-1.)
gen_wminus_decaytype[ak.num(gen_charged_leps_fromWminus) > 0] = np.ones(
ak.sum(ak.num(gen_charged_leps_fromWminus) > 0))
# add decaytype (0 is INVALID, 1 is LEPTONIC, 2 is HADRONIC)
gen_wminus['decaytype'] = ak.unflatten(gen_wminus_decaytype,
ak.num(gen_wminus))
gen_tbars['decaytype'] = gen_wminus['decaytype'] # set decaytype for tbars
gen_bbars['decaytype'] = gen_wminus['decaytype'] # set decaytype for bbars
# neutral leps from last W-
gen_neutral_leps_fromWminus = ak.flatten(last_gen_wminus.children[
(np.abs(last_gen_wminus.children.pdgId) == 12) |
(np.abs(last_gen_wminus.children.pdgId) == 14) |
(np.abs(last_gen_wminus.children.pdgId) == 16)],
axis=2)
gen_neutral_leps_fromWminus['charge'] = ak.zeros_like(
gen_neutral_leps_fromWminus.pt)
gen_wpartons_up = ak.Array({}, with_name="PtEtaPhiMLorentzVector")
gen_wpartons_dw = ak.Array({}, with_name="PtEtaPhiMLorentzVector")
gen_charged_leps = ak.Array({}, with_name="PtEtaPhiMLorentzVector")
gen_neutral_leps = ak.Array({}, with_name="PtEtaPhiMLorentzVector")
for column in columns:
if column == 'decaytype': continue
gen_wpartons_up[column] = ak.flatten(
ak.concatenate([
gen_wpartons_up_fromWplus[column],
gen_wpartons_up_fromWminus[column]
],
axis=1)) # (up-type partons from W+, W-)
gen_wpartons_dw[column] = ak.flatten(
ak.concatenate([
gen_wpartons_dw_fromWplus[column],
gen_wpartons_dw_fromWminus[column]
],
axis=1)) # (dw-type partons from W+, W-)
gen_charged_leps[column] = ak.flatten(
ak.concatenate([
gen_charged_leps_fromWplus[column],
gen_charged_leps_fromWminus[column]
],
axis=1)) # (charged leps from W+, W-)
gen_neutral_leps[column] = ak.flatten(
ak.concatenate([
gen_neutral_leps_fromWplus[column],
gen_neutral_leps_fromWminus[column]
],
axis=1)) # (neutral leps from W+, W-)
gen_wpartons_up = ak.unflatten(
gen_wpartons_up,
ak.num(gen_wpartons_up_fromWplus) + ak.num(gen_wpartons_up_fromWminus))
gen_wpartons_dw = ak.unflatten(
gen_wpartons_dw,
ak.num(gen_wpartons_dw_fromWplus) + ak.num(gen_wpartons_dw_fromWminus))
gen_charged_leps = ak.unflatten(
gen_charged_leps,
ak.num(gen_charged_leps_fromWplus) +
ak.num(gen_charged_leps_fromWminus))
gen_neutral_leps = ak.unflatten(
gen_neutral_leps,
ak.num(gen_neutral_leps_fromWplus) +
ak.num(gen_neutral_leps_fromWminus))
gen_taus = gen_charged_leps[np.abs(gen_charged_leps.pdgId) == 15]
gen_nu_taus = gen_neutral_leps[np.abs(gen_neutral_leps.pdgId) == 16]
# fully hadronic evts
had_evts = (ak.num(gen_charged_leps) == 0) & (
ak.num(gen_neutral_leps)
== 0) & (ak.num(gen_wpartons_up) == 2) & (ak.num(gen_wpartons_dw) == 2)
#set_trace()
# get direct tau decay products from hard processes (subset of gen_taus events above)
tau_decay_prods = genparts[genparts.hasFlags(
['isDirectHardProcessTauDecayProduct'])]
# only need decays to leptons (e/mu, tau nu) for event classification
tau_TO_tau_nu = tau_decay_prods[np.abs(tau_decay_prods.pdgId) == 16]
tau_TO_charged_lep = tau_decay_prods[(np.abs(tau_decay_prods.pdgId) == 11)
|
(np.abs(tau_decay_prods.pdgId) == 13)]
tau_TO_neutral_lep = tau_decay_prods[(np.abs(tau_decay_prods.pdgId) == 12)
|
(np.abs(tau_decay_prods.pdgId) == 14)]
# set decaytype for gen taus
charged_lep_decaytype_array = ak.to_numpy(
ak.flatten(ak.zeros_like(gen_charged_leps['pt'])))
# semilep evts
semilep_evts = (ak.num(gen_charged_leps) == 1) & (
ak.num(gen_neutral_leps)
== 1) & (ak.num(gen_wpartons_up) == 1) & (ak.num(gen_wpartons_dw) == 1)
tau_jets_evts = semilep_evts & (ak.num(gen_taus) == 1)
sl_evts_mask = np.repeat(ak.to_numpy(semilep_evts),
ak.to_numpy(ak.num(gen_charged_leps)))
semilep_decaytype_array = np.zeros(ak.to_numpy(semilep_evts).sum(),
dtype=int)
# tau -> l
semilep_tau_leptonic_decay = (tau_jets_evts) & (
ak.num(tau_TO_charged_lep) == 1) & (
ak.num(tau_TO_neutral_lep) == 1) & (ak.num(tau_TO_tau_nu) == 1)
semilep_tau_hadronic_decay = (tau_jets_evts) & (
~semilep_tau_leptonic_decay)
semilep_decaytype_array[
semilep_tau_leptonic_decay[semilep_evts]] = np.ones(
ak.to_numpy(semilep_tau_leptonic_decay).sum(), dtype=int)
semilep_decaytype_array[
semilep_tau_hadronic_decay[semilep_evts]] = np.ones(
ak.to_numpy(semilep_tau_hadronic_decay).sum(), dtype=int) * 2
# set charged_lep_decatype_array for semileptonic events
charged_lep_decaytype_array[sl_evts_mask] = semilep_decaytype_array
# dilep evts
dilep_evts = (ak.num(gen_charged_leps) == 2) & (
ak.num(gen_neutral_leps)
== 2) & (ak.num(gen_wpartons_up) == 0) & (ak.num(gen_wpartons_dw) == 0)
lep_lep_evts = dilep_evts & (ak.num(gen_taus) == 0)
lep_tau_evts = dilep_evts & (ak.num(gen_taus) == 1)
tau_tau_evts = dilep_evts & (ak.num(gen_taus) == 2)
dl_evts_mask = np.repeat(ak.to_numpy(dilep_evts),
ak.to_numpy(ak.num(gen_charged_leps)))
dilep_decaytype_array = np.zeros((ak.to_numpy(dilep_evts).sum(), 2),
dtype=int)
# tau + tau
# tau + tau -> ll
dilep_TauTau_ll_decay = (tau_tau_evts) & (
ak.num(tau_TO_charged_lep) == 2) & (
ak.num(tau_TO_neutral_lep) == 2) & (ak.num(tau_TO_tau_nu) == 2)
dilep_decaytype_array[dilep_TauTau_ll_decay[dilep_evts]] = np.ones(
(ak.to_numpy(dilep_TauTau_ll_decay).sum(), 2), dtype=int)
# tau + tau -> hh
dilep_TauTau_hh_decay = (tau_tau_evts) & (
(ak.num(tau_TO_charged_lep) + ak.num(tau_TO_neutral_lep))
== 0) & (ak.num(tau_TO_tau_nu) == 2)
dilep_decaytype_array[dilep_TauTau_hh_decay[dilep_evts]] = np.ones(
(ak.to_numpy(dilep_TauTau_hh_decay).sum(), 2), dtype=int) * 2
# tau + tau -> lh
dilep_TauTau_lh_decay = (
tau_tau_evts) & ~(dilep_TauTau_ll_decay | dilep_TauTau_hh_decay)
# set index corresponding to leptonically decaying tau to 1, default array is set to 2
dl_TauTau_to_lh_decaytype_array = np.ones(
ak.to_numpy(dilep_TauTau_lh_decay).sum() * 2, dtype=int) * 2
lep_tau_mask = (np.repeat(
ak.to_numpy(
ak.flatten(
tau_TO_charged_lep[dilep_TauTau_lh_decay].parent.pdgId)),
2) == ak.flatten(gen_charged_leps[dilep_TauTau_lh_decay].pdgId))
dl_TauTau_to_lh_decaytype_array[lep_tau_mask] = np.ones(lep_tau_mask.sum(),
dtype=int)
dilep_decaytype_array[dilep_TauTau_lh_decay[
dilep_evts]] = dl_TauTau_to_lh_decaytype_array.reshape(
ak.to_numpy(dilep_TauTau_lh_decay).sum(), 2)
# lep + tau
# tau -> l
dilep_LepTau_l_decay = (lep_tau_evts) & (
ak.num(tau_TO_charged_lep) == 1) & (
ak.num(tau_TO_neutral_lep) == 1) & (ak.num(tau_TO_tau_nu) == 1)
# set index corresponding to tau to 1
dl_LepTau_to_Lep_decaytype_array = np.zeros(
ak.to_numpy(dilep_LepTau_l_decay).sum() * 2, dtype=int)
dl_LepTau_to_Lep_decaytype_array[ak.flatten(
np.abs(gen_charged_leps[dilep_LepTau_l_decay].pdgId) == 15)] = np.ones(
ak.sum(dilep_LepTau_l_decay), dtype=int)
dilep_decaytype_array[dilep_LepTau_l_decay[
dilep_evts]] = dl_LepTau_to_Lep_decaytype_array.reshape(
ak.sum(dilep_LepTau_l_decay), 2)
# tau -> h
dilep_LepTau_h_decay = (lep_tau_evts) & ~(dilep_LepTau_l_decay)
# set index corresponding to tau to 2
dl_LepTau_to_Had_decaytype_array = np.zeros(ak.sum(dilep_LepTau_h_decay) *
2,
dtype=int)
dl_LepTau_to_Had_decaytype_array[ak.flatten(
np.abs(gen_charged_leps[dilep_LepTau_h_decay].pdgId) == 15)] = np.ones(
ak.sum(dilep_LepTau_h_decay), dtype=int) * 2
dilep_decaytype_array[dilep_LepTau_h_decay[
dilep_evts]] = dl_LepTau_to_Had_decaytype_array.reshape(
ak.sum(dilep_LepTau_h_decay), 2)
# set charged_lep_decatype_array for dileptonic events
charged_lep_decaytype_array[dl_evts_mask] = dilep_decaytype_array.flatten()
# set charged lepton decaytype (defined for taus only, e/mu are 0) (1 is LEPTONIC, 2 is HADRONIC)
gen_charged_leps['decaytype'] = ak.unflatten(charged_lep_decaytype_array,
ak.num(gen_charged_leps))
# make awkward arrays of (top decay prods, tbar decay prods)
Gen_Top_Pairs = ak.Array({}, with_name="PtEtaPhiMLorentzVector")
Gen_B_Pairs = ak.Array({}, with_name="PtEtaPhiMLorentzVector")
Gen_W_Pairs = ak.Array({}, with_name="PtEtaPhiMLorentzVector")
Gen_Wparton_Pairs = ak.Array({}, with_name="PtEtaPhiMLorentzVector")
for column in columns:
Gen_Top_Pairs[column] = ak.flatten(
ak.concatenate([gen_tops[column], gen_tbars[column]],
axis=1)) # (top, tbar)
Gen_B_Pairs[column] = ak.flatten(
ak.concatenate([gen_bs[column], gen_bbars[column]],
axis=1)) # (b, bbar)
Gen_W_Pairs[column] = ak.flatten(
ak.concatenate([gen_wplus[column], gen_wminus[column]],
axis=1)) # (W+, W-)
if column is not "decaytype":
Gen_Wparton_Pairs[column] = ak.flatten(
ak.concatenate(
[
ak.pad_none(gen_wpartons_up[column], 1, axis=1),
ak.pad_none(gen_wpartons_dw[column], 1, axis=1)
],
axis=1)) # (up-type wpartons, down-type wpartons)
Gen_Top_Pairs = ak.unflatten(Gen_Top_Pairs,
ak.num(gen_tops) + ak.num(gen_tbars))
Gen_B_Pairs = ak.unflatten(Gen_B_Pairs, ak.num(gen_bs) + ak.num(gen_bbars))
Gen_W_Pairs = ak.unflatten(Gen_W_Pairs,
ak.num(gen_wplus) + ak.num(gen_wminus))
Gen_Wparton_Pairs = ak.unflatten(
Gen_Wparton_Pairs,
ak.num(ak.pad_none(gen_wpartons_up, 1, axis=1)) +
ak.num(ak.pad_none(gen_wpartons_dw, 1, axis=1)))
Gen_Wparton_Pairs = Gen_Wparton_Pairs[ak.argsort(
Gen_Wparton_Pairs["pt"], ascending=False)] # sort by pt
Gen_TTbar = ak.Array(
{
"pt": (gen_tops + gen_tbars).pt,
"eta": (gen_tops + gen_tbars).eta,
"phi": (gen_tops + gen_tbars).phi,
"mass": (gen_tops + gen_tbars).mass,
"decaytype": gen_tops["decaytype"] + gen_tbars[
"decaytype"], # 0 is for INVALID, 2 for DILEP, 3 for SEMILEP, 4 for HADRONIC
},
with_name="PtEtaPhiMLorentzVector")
## make "table" of gen objects for certain decays
## DILEP
DILEP_evts = ak.zip({
"TTbar":
Gen_TTbar[dilep_evts] if ak.any(dilep_evts) else ak.unflatten(
Gen_TTbar[dilep_evts], ak.values_astype(dilep_evts, int)),
"Top":
Gen_Top_Pairs[np.sign(Gen_Top_Pairs.charge) == 1][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
Gen_Top_Pairs[np.sign(Gen_Top_Pairs.charge) == 1][dilep_evts],
ak.values_astype(dilep_evts, int)),
"Tbar":
Gen_Top_Pairs[np.sign(Gen_Top_Pairs.charge) == -1][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
Gen_Top_Pairs[np.sign(Gen_Top_Pairs.charge) == -1][dilep_evts],
ak.values_astype(dilep_evts, int)),
"B":
Gen_B_Pairs[np.sign(Gen_B_Pairs.charge) == -1][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
Gen_B_Pairs[np.sign(Gen_B_Pairs.charge) == -1][dilep_evts],
ak.values_astype(dilep_evts, int)),
"Bbar":
Gen_B_Pairs[np.sign(Gen_B_Pairs.charge) == 1][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
Gen_B_Pairs[np.sign(Gen_B_Pairs.charge) == 1][dilep_evts],
ak.values_astype(dilep_evts, int)),
"Wplus":
Gen_W_Pairs[Gen_W_Pairs.charge == 1][dilep_evts] if ak.any(dilep_evts)
else ak.unflatten(Gen_W_Pairs[Gen_W_Pairs.charge == 1][dilep_evts],
ak.values_astype(dilep_evts, int)),
"Wminus":
Gen_W_Pairs[Gen_W_Pairs.charge == -1][dilep_evts] if ak.any(dilep_evts)
else ak.unflatten(Gen_W_Pairs[Gen_W_Pairs.charge == -1][dilep_evts],
ak.values_astype(dilep_evts, int)),
"First_plus":
gen_charged_leps[gen_charged_leps.charge > 0][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
gen_charged_leps[gen_charged_leps.charge > 0][dilep_evts],
ak.values_astype(dilep_evts,
int)), # charged lepton always made leading
"Second_plus":
gen_neutral_leps[gen_charged_leps.charge > 0][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
gen_neutral_leps[gen_charged_leps.charge > 0][dilep_evts],
ak.values_astype(dilep_evts,
int)), # neutral lepton always made subleading
"First_minus":
gen_charged_leps[gen_charged_leps.charge < 0][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
gen_charged_leps[gen_charged_leps.charge < 0][dilep_evts],
ak.values_astype(dilep_evts,
int)), # charged lepton always made leading
"Second_minus":
gen_neutral_leps[gen_charged_leps.charge < 0][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
gen_neutral_leps[gen_charged_leps.charge < 0][dilep_evts],
ak.values_astype(dilep_evts,
int)), # neutral lepton always made subleading
"Up_plus":
gen_neutral_leps[gen_charged_leps.charge > 0][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
gen_neutral_leps[gen_charged_leps.charge > 0][dilep_evts],
ak.values_astype(dilep_evts, int)), # same as second plus
"Down_plus":
gen_charged_leps[gen_charged_leps.charge > 0][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
gen_charged_leps[gen_charged_leps.charge > 0][dilep_evts],
ak.values_astype(dilep_evts, int)), # same as first plus
"Up_minus":
gen_neutral_leps[gen_charged_leps.charge < 0][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
gen_neutral_leps[gen_charged_leps.charge < 0][dilep_evts],
ak.values_astype(dilep_evts, int)), # same as second minus
"Down_minus":
gen_charged_leps[gen_charged_leps.charge < 0][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
gen_charged_leps[gen_charged_leps.charge < 0][dilep_evts],
ak.values_astype(dilep_evts, int)), # same as first minus
})
## HAD
HAD_evts = ak.zip({
"TTbar":
Gen_TTbar[had_evts] if ak.any(had_evts) else ak.unflatten(
Gen_TTbar[had_evts], ak.values_astype(had_evts, int)),
"Top":
Gen_Top_Pairs[np.sign(Gen_Top_Pairs.charge) == 1][had_evts]
if ak.any(had_evts) else ak.unflatten(
Gen_Top_Pairs[np.sign(Gen_Top_Pairs.charge) == 1][had_evts],
ak.values_astype(had_evts, int)),
"Tbar":
Gen_Top_Pairs[np.sign(Gen_Top_Pairs.charge) == -1][had_evts]
if ak.any(had_evts) else ak.unflatten(
Gen_Top_Pairs[np.sign(Gen_Top_Pairs.charge) == -1][had_evts],
ak.values_astype(had_evts, int)),
"B":
Gen_B_Pairs[np.sign(Gen_B_Pairs.charge) == -1][had_evts]
if ak.any(had_evts) else ak.unflatten(
Gen_B_Pairs[np.sign(Gen_B_Pairs.charge) == -1][had_evts],
ak.values_astype(had_evts, int)),
"Bbar":
Gen_B_Pairs[np.sign(Gen_B_Pairs.charge) == 1][had_evts]
if ak.any(had_evts) else ak.unflatten(
Gen_B_Pairs[np.sign(Gen_B_Pairs.charge) == 1][had_evts],
ak.values_astype(had_evts, int)),
"Wplus":
Gen_W_Pairs[Gen_W_Pairs.charge == 1][had_evts] if ak.any(had_evts) else
ak.unflatten(Gen_W_Pairs[Gen_W_Pairs.charge == 1][had_evts],
ak.values_astype(had_evts, int)),
"Wminus":
Gen_W_Pairs[Gen_W_Pairs.charge == -1][had_evts] if ak.any(had_evts)
else ak.unflatten(Gen_W_Pairs[Gen_W_Pairs.charge == -1][had_evts],
ak.values_astype(had_evts, int)),
"First_plus":
Gen_Wparton_Pairs[had_evts][Gen_Wparton_Pairs[had_evts].charge > 0][:,
0]
if ak.any(had_evts) else ak.unflatten(
Gen_Wparton_Pairs[had_evts][Gen_Wparton_Pairs[had_evts].charge > 0]
[:, 0], ak.values_astype(
had_evts, int)), # leading positively-charged parton
"Second_plus":
Gen_Wparton_Pairs[had_evts][Gen_Wparton_Pairs[had_evts].charge > 0][:,
1]
if ak.any(had_evts) else ak.unflatten(
Gen_Wparton_Pairs[had_evts][Gen_Wparton_Pairs[had_evts].charge > 0]
[:, 1], ak.values_astype(
had_evts, int)), # subleading positively-charged parton
"First_minus":
Gen_Wparton_Pairs[had_evts][Gen_Wparton_Pairs[had_evts].charge < 0][:,
0]
if ak.any(had_evts) else ak.unflatten(
Gen_Wparton_Pairs[had_evts][Gen_Wparton_Pairs[had_evts].charge < 0]
[:, 0], ak.values_astype(
had_evts, int)), # leading negatively-charged parton
"Second_minus":
Gen_Wparton_Pairs[had_evts][Gen_Wparton_Pairs[had_evts].charge < 0][:,
1]
if ak.any(had_evts) else ak.unflatten(
Gen_Wparton_Pairs[had_evts][Gen_Wparton_Pairs[had_evts].charge < 0]
[:, 1], ak.values_astype(
had_evts, int)), # subleading negatively-charged parton
"Up_plus":
gen_wpartons_up[gen_wpartons_up.charge > 0][had_evts]
if ak.any(had_evts) else ak.unflatten(
gen_wpartons_up[gen_wpartons_up.charge > 0][had_evts],
ak.values_astype(had_evts,
int)), # positively-charged up-type parton
"Down_plus":
gen_wpartons_dw[gen_wpartons_dw.charge > 0][had_evts]
if ak.any(had_evts) else ak.unflatten(
gen_wpartons_dw[gen_wpartons_dw.charge > 0][had_evts],
ak.values_astype(had_evts,
int)), # positively-charged down-type parton
"Up_minus":
gen_wpartons_up[gen_wpartons_up.charge < 0][had_evts]
if ak.any(had_evts) else ak.unflatten(
gen_wpartons_up[gen_wpartons_up.charge < 0][had_evts],
ak.values_astype(had_evts,
int)), # negatively-charged up-type parton
"Down_minus":
gen_wpartons_dw[gen_wpartons_dw.charge < 0][had_evts]
if ak.any(had_evts) else ak.unflatten(
gen_wpartons_dw[gen_wpartons_dw.charge < 0][had_evts],
ak.values_astype(had_evts,
int)), # negatively-charged down-type parton
})
## SEMILEP
SEMILEP_evts = ak.zip({
"TTbar":
Gen_TTbar[semilep_evts] if ak.any(semilep_evts) else ak.unflatten(
Gen_TTbar[semilep_evts], ak.values_astype(semilep_evts, int)),
"THad":
Gen_Top_Pairs[Gen_Top_Pairs.decaytype == 2][semilep_evts]
if ak.any(semilep_evts) else ak.unflatten(
Gen_Top_Pairs[Gen_Top_Pairs.decaytype == 2][semilep_evts],
ak.values_astype(semilep_evts, int)),
"TLep":
Gen_Top_Pairs[Gen_Top_Pairs.decaytype == 1][semilep_evts]
if ak.any(semilep_evts) else ak.unflatten(
Gen_Top_Pairs[Gen_Top_Pairs.decaytype == 1][semilep_evts],
ak.values_astype(semilep_evts, int)),
"BHad":
Gen_B_Pairs[Gen_B_Pairs.decaytype == 2][semilep_evts]
if ak.any(semilep_evts) else ak.unflatten(
Gen_B_Pairs[Gen_B_Pairs.decaytype == 2][semilep_evts],
ak.values_astype(semilep_evts, int)),
"BLep":
Gen_B_Pairs[Gen_B_Pairs.decaytype == 1][semilep_evts]
if ak.any(semilep_evts) else ak.unflatten(
Gen_B_Pairs[Gen_B_Pairs.decaytype == 1][semilep_evts],
ak.values_astype(semilep_evts, int)),
"WHad":
Gen_W_Pairs[Gen_W_Pairs.decaytype == 2][semilep_evts]
if ak.any(semilep_evts) else ak.unflatten(
Gen_W_Pairs[Gen_W_Pairs.decaytype == 2][semilep_evts],
ak.values_astype(semilep_evts, int)),
"WLep":
Gen_W_Pairs[Gen_W_Pairs.decaytype == 1][semilep_evts]
if ak.any(semilep_evts) else ak.unflatten(
Gen_W_Pairs[Gen_W_Pairs.decaytype == 1][semilep_evts],
ak.values_astype(semilep_evts, int)),
"Lepton":
gen_charged_leps[semilep_evts] if ak.any(semilep_evts) else
ak.unflatten(gen_charged_leps[semilep_evts],
ak.values_astype(semilep_evts, int)),
"Nu":
gen_neutral_leps[semilep_evts] if ak.any(semilep_evts) else
ak.unflatten(gen_neutral_leps[semilep_evts],
ak.values_astype(semilep_evts, int)),
"WJa":
Gen_Wparton_Pairs[:, 0][semilep_evts] if ak.any(semilep_evts) else
ak.unflatten(Gen_Wparton_Pairs[:, 0][semilep_evts],
ak.values_astype(semilep_evts, int)),
"WJb":
Gen_Wparton_Pairs[:, 1][semilep_evts] if ak.any(semilep_evts) else
ak.unflatten(Gen_Wparton_Pairs[:, 1][semilep_evts],
ak.values_astype(semilep_evts, int)),
"Up_Had":
gen_wpartons_up[semilep_evts] if ak.any(semilep_evts) else
ak.unflatten(gen_wpartons_up[semilep_evts],
ak.values_astype(semilep_evts, int)),
"Down_Had":
gen_wpartons_dw[semilep_evts] if ak.any(semilep_evts) else
ak.unflatten(gen_wpartons_dw[semilep_evts],
ak.values_astype(semilep_evts, int)),
})
# make dictionary to return
GenObjects = dict({
"SL": SEMILEP_evts,
"DL": DILEP_evts,
"Had": HAD_evts,
})
return GenObjects
def test_rochester():
rochester_data = lookup_tools.txt_converters.convert_rochester_file(
"tests/samples/RoccoR2018.txt.gz", loaduncs=True)
rochester = lookup_tools.rochester_lookup.rochester_lookup(rochester_data)
# to test 1-to-1 agreement with official Rochester requires loading C++ files
# instead, preload the correct scales in the sample directory
# the script tests/samples/rochester/build_rochester.py produces these
official_data_k = np.load("tests/samples/nano_dimuon_rochester.npy")
official_data_err = np.load("tests/samples/nano_dimuon_rochester_err.npy")
official_mc_k = np.load("tests/samples/nano_dy_rochester.npy")
official_mc_err = np.load("tests/samples/nano_dy_rochester_err.npy")
mc_rand = np.load("tests/samples/nano_dy_rochester_rand.npy")
# test against nanoaod
events = NanoEventsFactory.from_root(
os.path.abspath("tests/samples/nano_dimuon.root")).events()
data_k = rochester.kScaleDT(events.Muon.charge, events.Muon.pt,
events.Muon.eta, events.Muon.phi)
data_k = np.array(ak.flatten(data_k))
assert all(np.isclose(data_k, official_data_k))
data_err = rochester.kScaleDTerror(events.Muon.charge, events.Muon.pt,
events.Muon.eta, events.Muon.phi)
data_err = np.array(ak.flatten(data_err), dtype=float)
assert all(np.isclose(data_err, official_data_err, atol=1e-8))
# test against mc
events = NanoEventsFactory.from_root(
os.path.abspath("tests/samples/nano_dy.root")).events()
hasgen = ~np.isnan(ak.fill_none(events.Muon.matched_gen.pt, np.nan))
mc_rand = ak.unflatten(mc_rand, ak.num(hasgen))
mc_kspread = rochester.kSpreadMC(
events.Muon.charge[hasgen],
events.Muon.pt[hasgen],
events.Muon.eta[hasgen],
events.Muon.phi[hasgen],
events.Muon.matched_gen.pt[hasgen],
)
mc_ksmear = rochester.kSmearMC(
events.Muon.charge[~hasgen],
events.Muon.pt[~hasgen],
events.Muon.eta[~hasgen],
events.Muon.phi[~hasgen],
events.Muon.nTrackerLayers[~hasgen],
mc_rand[~hasgen],
)
mc_k = np.array(ak.flatten(ak.ones_like(events.Muon.pt)))
hasgen_flat = np.array(ak.flatten(hasgen))
mc_k[hasgen_flat] = np.array(ak.flatten(mc_kspread))
mc_k[~hasgen_flat] = np.array(ak.flatten(mc_ksmear))
assert all(np.isclose(mc_k, official_mc_k))
mc_errspread = rochester.kSpreadMCerror(
events.Muon.charge[hasgen],
events.Muon.pt[hasgen],
events.Muon.eta[hasgen],
events.Muon.phi[hasgen],
events.Muon.matched_gen.pt[hasgen],
)
mc_errsmear = rochester.kSmearMCerror(
events.Muon.charge[~hasgen],
events.Muon.pt[~hasgen],
events.Muon.eta[~hasgen],
events.Muon.phi[~hasgen],
events.Muon.nTrackerLayers[~hasgen],
mc_rand[~hasgen],
)
mc_err = np.array(ak.flatten(ak.ones_like(events.Muon.pt)))
mc_err[hasgen_flat] = np.array(ak.flatten(mc_errspread))
mc_err[~hasgen_flat] = np.array(ak.flatten(mc_errsmear))
assert all(np.isclose(mc_err, official_mc_err, atol=1e-8))
def build(self, jets, lazy_cache):
if lazy_cache is None:
raise Exception(
'CorrectedJetsFactory requires a awkward-array cache to function correctly.'
)
fields = None
out = None
wrap = None
VirtualType = None
form = None
if isinstance(jets, awkward.highlevel.Array):
fields = awkward.fields(jets)
if len(fields) == 0:
raise Exception(
'Detected awkward: \'jets\' must have attributes specified by keys!'
)
out = awkward.flatten(jets)
wrap = partial(awkward_rewrap,
like_what=jets,
gfunc=rewrap_recordarray)
VirtualType = awkward.virtual
form = out[self.name_map['ptRaw']].layout.form
else:
raise Exception(
'\'jets\' must be an awkward > 1.0.0 array of some kind!')
if len(fields) == 0:
raise Exception(
'Empty record, please pass a jet object with at least {self.real_sig} defined!'
)
# take care of nominal JEC (no JER if available)
out[self.name_map['JetPt'] + '_orig'] = out[self.name_map['JetPt']]
out[self.name_map['JetMass'] + '_orig'] = out[self.name_map['JetMass']]
if self.treat_pt_as_raw:
out[self.name_map['ptRaw']] = out[self.name_map['JetPt']]
out[self.name_map['massRaw']] = out[self.name_map['JetMass']]
jec_name_map = {}
jec_name_map.update(self.name_map)
jec_name_map['JetPt'] = jec_name_map['ptRaw']
jec_name_map['JetMass'] = jec_name_map['ptRaw']
if self.jec_stack.jec is not None:
jec_args = {
k: out[jec_name_map[k]]
for k in self.jec_stack.jec.signature
}
out['jet_energy_correction'] = self.jec_stack.jec.getCorrection(
**jec_args, form=form, lazy_cache=lazy_cache)
else:
out['jet_energy_correction'] = awkward.ones_like(
out[self.name_map['JetPt']])
# finally the lazy binding to the JEC
def jec_var_corr(arr, varName):
return arr['jet_energy_correction'] * arr[varName]
init_pt = partial(VirtualType,
jec_var_corr,
args=(out, self.name_map['ptRaw']),
cache=lazy_cache)
init_mass = partial(VirtualType,
jec_var_corr,
args=(out, self.name_map['massRaw']),
cache=lazy_cache)
out[self.name_map['JetPt']] = init_pt(length=len(out), form=form)
out[self.name_map['JetMass']] = init_mass(length=len(out), form=form)
out[self.name_map['JetPt'] + '_jec'] = init_pt(length=len(out),
form=form)
out[self.name_map['JetMass'] + '_jec'] = init_mass(length=len(out),
form=form)
# in jer we need to have a stash for the intermediate JEC products
has_jer = False
if self.jec_stack.jer is not None and self.jec_stack.jersf is not None:
has_jer = True
jer_name_map = {}
jer_name_map.update(self.name_map)
jer_name_map['JetPt'] = jer_name_map['JetPt'] + '_jec'
jer_name_map['JetMass'] = jer_name_map['JetMass'] + '_jec'
jerargs = {
k: out[jer_name_map[k]]
for k in self.jec_stack.jer.signature
}
out['jet_energy_resolution'] = self.jec_stack.jer.getResolution(
**jerargs, form=form, lazy_cache=lazy_cache)
jersfargs = {
k: out[jer_name_map[k]]
for k in self.jec_stack.jersf.signature
}
out['jet_energy_resolution_scale_factor'] = self.jec_stack.jersf.getScaleFactor(
**jersfargs, form=_JERSF_FORM, lazy_cache=lazy_cache)
init_rand_gauss = partial(VirtualType,
args=(out, jer_name_map['JetPt']),
cache=lazy_cache)
out['jet_resolution_rand_gauss'] = init_rand_gauss(rand_gauss,
length=len(out),
form=form)
init_jerc = partial(
VirtualType,
jer_smear,
args=(out, 0, self.forceStochastic, jer_name_map['ptGenJet'],
jer_name_map['JetPt'], jer_name_map['JetEta']),
cache=lazy_cache)
out['jet_energy_resolution_correction'] = init_jerc(
length=len(out), form=form)
def jer_smeared_val(arr, base, varName):
return arr['jet_energy_resolution_correction'] * base[varName]
init_pt_jer = partial(VirtualType,
jer_smeared_val,
args=(out, out, jer_name_map['JetPt']),
cache=lazy_cache)
init_mass_jer = partial(VirtualType,
jer_smeared_val,
args=(out, out, jer_name_map['JetPt']),
cache=lazy_cache)
out[self.name_map['JetPt']] = init_pt_jer(length=len(out),
form=form)
out[self.name_map['JetMass']] = init_mass_jer(length=len(out),
form=form)
out[self.name_map['JetPt'] + '_jer'] = init_pt_jer(length=len(out),
form=form)
out[self.name_map['JetMass'] + '_jer'] = init_mass_jer(
length=len(out), form=form)
# JER systematics
jerc_up = partial(
VirtualType,
jer_smear,
args=(out, 1, self.forceStochastic, jer_name_map['ptGenJet'],
jer_name_map['JetPt'], jer_name_map['JetEta']),
cache=lazy_cache)
up = awkward.flatten(jets)
up['jet_energy_resolution_correction'] = jerc_up(length=len(out),
form=form)
init_pt_jer = partial(VirtualType,
jer_smeared_val,
args=(up, out, jer_name_map['JetPt']),
cache=lazy_cache)
init_mass_jer = partial(VirtualType,
jer_smeared_val,
args=(up, out, jer_name_map['JetPt']),
cache=lazy_cache)
up[self.name_map['JetPt']] = init_pt_jer(length=len(out),
form=form)
up[self.name_map['JetMass']] = init_mass_jer(length=len(out),
form=form)
jerc_down = partial(
VirtualType,
jer_smear,
args=(out, 2, self.forceStochastic, jer_name_map['ptGenJet'],
jer_name_map['JetPt'], jer_name_map['JetEta']),
cache=lazy_cache)
down = awkward.flatten(jets)
down['jet_energy_resolution_correction'] = jerc_down(
length=len(out), form=form)
init_pt_jer = partial(VirtualType,
jer_smeared_val,
args=(down, out, jer_name_map['JetPt']),
cache=lazy_cache)
init_mass_jer = partial(VirtualType,
jer_smeared_val,
args=(down, out, jer_name_map['JetPt']),
cache=lazy_cache)
down[self.name_map['JetPt']] = init_pt_jer(length=len(out),
form=form)
down[self.name_map['JetMass']] = init_mass_jer(length=len(out),
form=form)
out['JER'] = awkward.zip({
'up': up,
'down': down
},
depth_limit=1,
with_name='JetSystematic')
if self.jec_stack.junc is not None:
juncnames = {}
juncnames.update(self.name_map)
if has_jer:
juncnames['JetPt'] = juncnames['JetPt'] + '_jer'
juncnames['JetMass'] = juncnames['JetMass'] + '_jer'
else:
juncnames['JetPt'] = juncnames['JetPt'] + '_jec'
juncnames['JetMass'] = juncnames['JetMass'] + '_jec'
juncargs = {
k: out[juncnames[k]]
for k in self.jec_stack.junc.signature
}
juncs = self.jec_stack.junc.getUncertainty(**juncargs)
for name, func in juncs:
out[f'jet_energy_uncertainty_{name}'] = func
def junc_smeared_val(arr, up_down, uncname, varName):
return arr[
f'jet_energy_uncertainty_{uncname}'][:, up_down] * arr[
varName]
up = awkward.flatten(jets)
up[self.name_map['JetPt']] = VirtualType(
junc_smeared_val,
args=(
out,
0,
name,
juncnames['JetPt'],
),
length=len(out),
form=form,
cache=lazy_cache)
up[self.name_map['JetMass']] = VirtualType(
junc_smeared_val,
args=(
out,
0,
name,
juncnames['JetMass'],
),
length=len(out),
form=form,
cache=lazy_cache)
down = awkward.flatten(jets)
down[self.name_map['JetPt']] = VirtualType(
junc_smeared_val,
args=(
out,
1,
name,
juncnames['JetPt'],
),
length=len(out),
form=form,
cache=lazy_cache)
down[self.name_map['JetMass']] = VirtualType(
junc_smeared_val,
args=(
out,
1,
name,
juncnames['JetMass'],
),
length=len(out),
form=form,
cache=lazy_cache)
out[f'JES_{name}'] = awkward.zip({
'up': up,
'down': down
},
depth_limit=1,
with_name='JetSystematic')
return wrap(out)
