def tth_leptonic_preselection(events, photons, electrons, muons, jets, options,
debug):
"""
Performs tth leptonic preselection, requiring >= 1 lepton and >= 1 jet
Assumes diphoton preselection has already been applied.
Also calculates relevant event-level variables.
"""
cut_diagnostics = utils.CutDiagnostics(
events=events,
debug=debug,
cut_set="[analysis_selections.py : tth_leptonic_preselection]")
# Get number of electrons, muons
selected_electrons = electrons[lepton_selections.select_electrons(
events, photons, electrons, options, debug)]
selected_muons = muons[lepton_selections.select_muons(
events, photons, muons, options, debug)]
n_electrons = awkward.num(selected_electrons)
n_muons = awkward.num(selected_muons)
n_leptons = n_electrons + n_muons
# Get number of jets
selected_jets = jets[jet_selections.select_jets(events, photons,
selected_electrons,
selected_muons, None, jets,
options, debug)]
n_jets = awkward.num(selected_jets)
lep_cut = n_leptons >= 1
jet_cut = n_jets >= 1
all_cuts = lep_cut & jet_cut
cut_diagnostics.add_cuts([lep_cut, jet_cut, all_cuts],
["N_leptons >= 1", "N_jets >= 1", "all"])
# Keep only selected events
selected_events = events[all_cuts]
selected_photons = photons[all_cuts]
selected_electrons = selected_electrons[all_cuts]
selected_muons = selected_muons[all_cuts]
selected_jets = selected_jets[all_cuts]
# Calculate event-level variables
selected_events = lepton_selections.set_electrons(selected_events,
selected_electrons,
debug)
selected_events = lepton_selections.set_muons(selected_events,
selected_muons, debug)
selected_events = jet_selections.set_jets(selected_events, selected_jets,
options, debug)
return selected_events
def test_nested_lorentz_vectorization(collection, events):
mask = ak.num(events[collection]) > 0
assert ak.all(ak.num(events[collection].PrunedP4_5, axis=2) == 5)
assert (ak.type(events[collection][mask].PrunedP4_5[
0, 0, 0]).parameters["__record__"] == "LorentzVector")
assert ak.all(ak.num(events[collection].SoftDroppedP4_5, axis=2) == 5)
assert (ak.type(events[collection][mask].SoftDroppedP4_5[
0, 0, 0]).parameters["__record__"] == "LorentzVector")
assert ak.all(ak.num(events[collection].TrimmedP4_5, axis=2) == 5)
assert (ak.type(events[collection][mask].TrimmedP4_5[
0, 0, 0]).parameters["__record__"] == "LorentzVector")
def diphoton_preselection(events, photons, options, debug):
# Initialize cut diagnostics tool for debugging
cut_diagnostics = utils.CutDiagnostics(
events=events,
debug=debug,
cut_set="[photon_selections.py : diphoton_preselection]")
selected_photons = photons[photon_selections.select_photons(
events, photons, options, debug)]
### mgg cut ###
resonant = options["resonant"]
if resonant:
mgg_mask = numpy.array(
events.ggMass > options["diphotons"]["mgg_lower"]) & numpy.array(
events.ggMass < options["diphotons"]["mgg_upper"])
else:
sideband_low = numpy.array(
events.ggMass > options["diphotons"]["mgg_lower"]) & numpy.array(
events.ggMass < options["diphotons"]["mgg_sideband_lower"])
sideband_high = numpy.array(
events.ggMass > options["diphotons"]["mgg_sideband_upper"]
) & numpy.array(events.ggMass < options["diphotons"]["mgg_upper"])
mgg_mask = sideband_low | sideband_high
### pt/mgg cuts ###
lead_pt_mgg_requirement = (selected_photons.pt / events.ggMass
) > options["photons"]["lead_pt_mgg_cut"]
sublead_pt_mgg_requirement = (selected_photons.pt / events.ggMass
) > options["photons"]["sublead_pt_mgg_cut"]
lead_pt_mgg_cut = awkward.num(
selected_photons[lead_pt_mgg_requirement]
) >= 1 # at least 1 photon passing lead requirement
sublead_pt_mgg_cut = awkward.num(
selected_photons[sublead_pt_mgg_requirement]
) >= 2 # at least 2 photon passing sublead requirement
pt_mgg_cut = lead_pt_mgg_cut & sublead_pt_mgg_cut
### 2 good selected_photons ###
n_photon_cut = awkward.num(
selected_photons
) == 2 # can regain a few % of signal if we set to >= 2 (probably e's that are reconstructed as selected_photons)
all_cuts = mgg_mask & pt_mgg_cut & n_photon_cut
cut_diagnostics.add_cuts([mgg_mask, pt_mgg_cut, n_photon_cut, all_cuts], [
"mgg in [100, 180]" if resonant else "mgg in [100, 120] or [130, 180]",
"lead (sublead) pt/mgg > 0.33 (0.25)", "2 good photons", "all"
])
return events[all_cuts], selected_photons[all_cuts]
def test_list_array():
array = ak.Array(np.arange(3 * 5 * 2).reshape(3, 5, 2).tolist())
assert ak.num(array, axis=0) == 3
assert ak.num(array, axis=1).tolist() == [5, 5, 5]
assert ak.num(array, axis=2).tolist() == [
[2, 2, 2, 2, 2],
[2, 2, 2, 2, 2],
[2, 2, 2, 2, 2],
]
with pytest.raises(ValueError) as err:
assert ak.num(array, axis=3)
assert str(err.value).startswith("'axis' out of range for 'num'")
assert ak.num(array, axis=-1).tolist() == [
[2, 2, 2, 2, 2],
[2, 2, 2, 2, 2],
[2, 2, 2, 2, 2],
]
assert ak.num(array, axis=-2).tolist() == [5, 5, 5]
assert ak.num(array, axis=-3) == 3
with pytest.raises(ValueError) as err:
assert ak.num(array, axis=-4)
assert str(err.value).startswith(
"axis == -4 exceeds the depth == 3 of this array")
def test_array_3d():
array = ak.Array(np.arange(3 * 5 * 2).reshape(3, 5, 2))
assert ak.to_list(array) == [
[[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]],
[[10, 11], [12, 13], [14, 15], [16, 17], [18, 19]],
[[20, 21], [22, 23], [24, 25], [26, 27], [28, 29]],
]
assert ak.num(array, axis=0) == 3
assert ak.to_list(ak.num(array, axis=1)) == [5, 5, 5]
assert ak.to_list(ak.num(array, axis=2)) == [
[2, 2, 2, 2, 2],
[2, 2, 2, 2, 2],
[2, 2, 2, 2, 2],
]
with pytest.raises(ValueError) as err:
assert ak.num(array, axis=3)
assert str(err.value).startswith("'axis' out of range for 'num'")
assert ak.to_list(ak.num(array, axis=-1)) == [
[2, 2, 2, 2, 2],
[2, 2, 2, 2, 2],
[2, 2, 2, 2, 2],
]
assert ak.to_list(ak.num(array, axis=-2)) == [5, 5, 5]
assert ak.num(array, axis=-3) == 3
with pytest.raises(ValueError) as err:
assert ak.num(array, axis=-4)
assert str(err.value).startswith(
"axis == -4 exceeds the depth == 3 of this array")
def awkward_to_wfsim_row_style(interactions):
"""
Converts awkward array instructions into instructions required by
WFSim.
:param interactions: awkward.Array containing GEANT4 simulation
information.
:return: Structured numpy.array. Each row represents either a S1 or
S2
"""
if len(interactions) == 0:
return np.empty(0, dtype=wfsim.instruction_dtype)
ninteractions = np.sum(ak.num(interactions['ed']))
res = np.zeros(2 * ninteractions, dtype=wfsim.instruction_dtype)
# TODO: Currently not supported rows with only electrons or photons due to
# this super odd shape
for i in range(2):
res['event_number'][i::2] = offset_range(
ak.to_numpy(ak.num(interactions['evtid'])))
res['type'][i::2] = i + 1
res['x'][i::2] = awkward_to_flat_numpy(interactions['x'])
res['y'][i::2] = awkward_to_flat_numpy(interactions['y'])
res['z'][i::2] = awkward_to_flat_numpy(interactions['z'])
res['x_pri'][i::2] = awkward_to_flat_numpy(interactions['x_pri'])
res['y_pri'][i::2] = awkward_to_flat_numpy(interactions['y_pri'])
res['z_pri'][i::2] = awkward_to_flat_numpy(interactions['z_pri'])
res['time'][i::2] = awkward_to_flat_numpy(interactions['t'])
res['g4id'][i::2] = awkward_to_flat_numpy(interactions['evtid'])
res['vol_id'][i::2] = awkward_to_flat_numpy(interactions['vol_id'])
res['e_dep'][i::2] = awkward_to_flat_numpy(interactions['ed'])
if 'local_field' in res.dtype.names:
res['local_field'][i::2] = awkward_to_flat_numpy(
interactions['e_field'])
recoil = awkward_to_flat_numpy(interactions['nestid'])
res['recoil'][i::2] = np.where(np.isin(recoil, [0, 6, 7, 8, 11]),
recoil, 8)
if i:
res['amp'][i::2] = awkward_to_flat_numpy(interactions['electrons'])
else:
res['amp'][i::2] = awkward_to_flat_numpy(interactions['photons'])
if 'n_excitons' in res.dtype.names:
res['n_excitons'][i::2] = awkward_to_flat_numpy(
interactions['excitons'])
# Remove entries with no quanta
res = res[res['amp'] > 0]
return res
def add_object_cuts(self, objects, object_name, cuts, names, sample):
if object_name not in self.object_stats.keys():
self.object_stats[object_name] = {}
if sample not in self.object_stats[object_name].keys():
self.object_stats[object_name][sample] = { "n_objects_initial" : awkward.sum(awkward.num(objects)) }
else:
self.object_stats[object_name][sample]["n_objects_initial"] += awkward.sum(awkward.num(objects))
for cut, name in zip(cuts, names):
if name not in self.object_stats[object_name][sample].keys():
self.object_stats[object_name][sample][name] = awkward.sum(awkward.num(objects[cut]))
else:
self.object_stats[object_name][sample][name] += awkward.sum(awkward.num(objects[cut]))
def runNN(model, varsIn, varSet, normMean, normStd):
dataset = RootDataset(varsIn=varsIn,
varSet=varSet,
normMean=normMean,
normStd=normStd)
nnOutput = getNNOutput(dataset, model)
fjets = varsIn["fjets"]
counts = ak.num(fjets.pt)
svjJetsAK8 = ak.unflatten(nnOutput, counts)
wpt = 0.5
darksvjJetsAK8 = fjets[svjJetsAK8 >= wpt]
varsIn['nsvjJetsAK8'] = [ak.num(darksvjJetsAK8), "evtw"]
varsIn['nnOutput'] = [svjJetsAK8, "fjw"]
def calculate_selection(self, syst_tag, events):
"""
"""
electron_cut = lepton_selections.select_electrons(
electrons=events.Electron,
options=self.options["electrons"],
clean={},
name="ele",
tagger=self)
electrons = awkward_utils.add_field(events=events,
name="ele",
data=events.Electron[electron_cut])
electrons = awkward.Array(electrons, with_name="Momentum4D")
ee_pairs = awkward.combinations(electrons,
2,
fields=["LeadEle", "SubleadEle"])
ee_pairs["ZCand"] = ee_pairs.LeadEle + ee_pairs.SubleadEle
ee_pairs[("ZCand", "mass")] = ee_pairs.ZCand.mass
ee_pairs[("ZCand", "pt")] = ee_pairs.ZCand.pt
ee_pairs[("ZCand", "phi")] = ee_pairs.ZCand.phi
ee_pairs[("ZCand", "eta")] = ee_pairs.ZCand.eta
events["ZCand"] = ee_pairs.ZCand
os_cut = ee_pairs.LeadEle.charge * ee_pairs.SubleadEle.charge == -1
mass_cut = (ee_pairs.ZCand.mass > self.options["z_window"][0]) & (
ee_pairs.ZCand.mass < self.options["z_window"][1])
id_cut = (ee_pairs.LeadEle.mvaFall17V2Iso_WP80
== True) | (ee_pairs.SubleadEle.mvaFall17V2Iso_WP80 == True)
pair_cut = os_cut & mass_cut & id_cut
evt_os_cut = awkward.num(ee_pairs[os_cut]) == 1
evt_mass_cut = awkward.num(ee_pairs[mass_cut]) == 1
evt_id_cut = awkward.num(ee_pairs[id_cut]) == 1
met_cut = events.MET_pt <= self.options["met"]
presel_cut = (awkward.num(ee_pairs[pair_cut]) == 1) & met_cut
self.register_cuts(
names=["met cut", "os cut", "mass cut", "id cut", "all"],
results=[
met_cut, evt_os_cut, evt_mass_cut, evt_id_cut, presel_cut
])
return presel_cut, events
def test_record_array():
array = ak.Array([
{
"x": [1],
"y": [[], [1]]
},
{
"x": [1, 2],
"y": [[], [1], [1, 2]]
},
{
"x": [1, 2, 3],
"y": [[], [1], [1, 2], [1, 2, 3]]
},
])
assert ak.num(array, axis=0).tolist() == {"x": 3, "y": 3}
assert ak.num(array, axis=1).tolist() == [
{
"x": 1,
"y": 2
},
{
"x": 2,
"y": 3
},
{
"x": 3,
"y": 4
},
]
with pytest.raises(ValueError) as err:
assert ak.num(array, axis=2)
assert str(err.value).startswith("'axis' out of range for 'num'")
assert ak.num(array, axis=-1).tolist() == [
{
"x": 1,
"y": [0, 1]
},
{
"x": 2,
"y": [0, 1, 2]
},
{
"x": 3,
"y": [0, 1, 2, 3]
},
]
def test():
array = ak.Array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
assert ak.unflatten(array, 5).tolist() == [[0, 1, 2, 3, 4],
[5, 6, 7, 8, 9]]
assert ak.unflatten(array, [3, 0, 2, 1, 4]).tolist() == [
[0, 1, 2],
[],
[3, 4],
[5],
[6, 7, 8, 9],
]
assert ak.unflatten(array, [3, None, 2, 1, 4]).tolist() == [
[0, 1, 2],
None,
[3, 4],
[5],
[6, 7, 8, 9],
]
original = ak.Array([[0, 1, 2], [], [3, 4], [5], [6, 7, 8, 9]])
counts = ak.num(original)
array = ak.flatten(original)
assert counts.tolist() == [3, 0, 2, 1, 4]
assert array.tolist() == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
assert ak.unflatten(array, counts).tolist() == [
[0, 1, 2],
[],
[3, 4],
[5],
[6, 7, 8, 9],
]
def set_fatjets(events, fatjets, options, debug):
events["n_fatjets"] = awkward.num(fatjets)
n_save = 1
fatjet_pt_padded = utils.pad_awkward_array(fatjets.pt, n_save, -9)
fatjet_eta_padded = utils.pad_awkward_array(fatjets.eta, n_save, -9)
fatjet_mass_padded = utils.pad_awkward_array(fatjets.mass, n_save, -9)
fatjet_msoftdrop_padded = utils.pad_awkward_array(fatjets.msoftdrop,
n_save, -9)
fatjet_btag_padded = utils.pad_awkward_array(fatjets.btagDDBvL_noMD,
n_save, -9)
fatjet_deepbtag_md_padded = utils.pad_awkward_array(
fatjets.deepTagMD_HbbvsQCD, n_save, -9)
for i in range(n_save):
events["fatjet%s_pt" % str(i + 1)] = fatjet_pt_padded[:, i]
events["fatjet%s_eta" % str(i + 1)] = fatjet_eta_padded[:, i]
events["fatjet%s_msoftdrop" % str(i + 1)] = fatjet_msoftdrop_padded[:,
i]
events["fatjet%s_mass" % str(i + 1)] = fatjet_mass_padded[:, i]
events["fatjet%s_btag" % str(i + 1)] = fatjet_btag_padded[:, i]
events["fatjet%s_deepbtag_md" %
str(i + 1)] = fatjet_deepbtag_md_padded[:, i]
return events
def pad_and_reshape_nested_arrays(self,
batch,
variable_type,
max_items=10):
"""
Function that acts on nested data to read relevant variables, pad, reshape and convert data from uproot into
rectilinear numpy arrays
:param batch: A dict of awkward arrays from uproot
:param variable_type: Variable type to be selected e.g. Tracks, Neutral PFO, Jets etc...
:param max_items: Maximum number of tracks/PFOs etc... to be associated to event
:return: a rectilinear numpy array of shape:
(num events in batch, number of variables belonging to variable type, max_items)
"""
variables = self._variables_dict[variable_type]
np_arrays = np.zeros((ak.num(batch[variables[0]],
axis=0), len(variables), max_items))
dummy_val = 0
thresh = 45
for i in range(0, len(variables)):
ak_arr = batch[variables[i]]
ak_arr = ak.pad_none(ak_arr, max_items, clip=True, axis=1)
arr = ak.to_numpy(abs(ak_arr)).filled(dummy_val)
np_arrays[:, i] = arr
np_arrays = apply_scaling(np_arrays,
thresh=thresh,
dummy_val=dummy_val)
np_arrays = np.nan_to_num(np_arrays, posinf=0, neginf=0,
copy=False).astype("float64")
return np_arrays
def find_permutations(jets, leptons, MET, btagWP):
'''
Inputs:
Jets, leptons, MET, and if jets pass btag WP
Returns:
List of (jet assignment ordering, associated neutrino solutions)
'''
jets_inputs = np.stack((ak.to_numpy(ak.flatten(jets.px)), ak.to_numpy(ak.flatten(jets.py)), ak.to_numpy(ak.flatten(jets.pz)), ak.to_numpy(ak.flatten(jets.energy)), ak.to_numpy(ak.flatten(jets[btagWP]))), axis=1).astype('float64') # one row has (px, py, pyz, E)
lepton_inputs = np.stack((ak.to_numpy(ak.flatten(leptons.px)), ak.to_numpy(ak.flatten(leptons.py)), ak.to_numpy(ak.flatten(leptons.pz)), ak.to_numpy(ak.flatten(leptons.energy))), axis=1).astype('float64') # one row has (px, py, pyz, E)
met_inputs = np.stack((ak.to_numpy(MET.px), ak.to_numpy(MET.py)), axis=1).astype('float64') # one row has (px, py)
p_ordering, p_nu = get_test_permutations(njets_array=ak.num(jets), jets=jets_inputs, leptons=lepton_inputs, met=met_inputs)
#set_trace()
test_perms = ak.Array({
'blepIdx' : ak.from_iter(p_ordering)[:, :, 0],
'bhadIdx' : ak.from_iter(p_ordering)[:, :, 1],
'wjaIdx' : ak.from_iter(p_ordering)[:, :, 2],
'wjbIdx' : ak.from_iter(p_ordering)[:, :, 3],
'Nu' : ak.Array({
'px' : ak.from_iter(p_nu)[:, :, 0],
'py' : ak.from_iter(p_nu)[:, :, 1],
'pz' : ak.from_iter(p_nu)[:, :, 2],
'chi2' : ak.from_iter(p_nu)[:, :, 3],
})
})
return test_perms
def set_taus(events, taus, debug):
events["n_tau"] = awkward.num(taus)
tau_pt_padded = utils.pad_awkward_array(taus.pt, 2, -9)
tau_eta_padded = utils.pad_awkward_array(taus.eta, 2, -9)
tau_phi_padded = utils.pad_awkward_array(taus.phi, 2, -9)
tau_mass_padded = utils.pad_awkward_array(taus.mass, 2, -9)
tau_IDvsElec_padded = utils.pad_awkward_array(taus.idDeepTau2017v2p1VSe, 2,
-9)
tau_IDvsJet_padded = utils.pad_awkward_array(taus.idDeepTau2017v2p1VSjet,
2, -9)
tau_IDvsMuon_padded = utils.pad_awkward_array(taus.idDeepTau2017v2p1VSmu,
2, -9)
events["tau1_pt"] = tau_pt_padded[:, 0]
events["tau2_pt"] = tau_pt_padded[:, 1]
events["tau1_eta"] = tau_eta_padded[:, 0]
events["tau2_eta"] = tau_eta_padded[:, 1]
events["tau1_phi"] = tau_phi_padded[:, 0]
events["tau2_phi"] = tau_phi_padded[:, 1]
events["tau1_mass"] = tau_mass_padded[:, 0]
events["tau2_mass"] = tau_mass_padded[:, 1]
events["tau1_id_vs_e"] = tau_IDvsElec_padded[:, 0]
events["tau2_id_vs_e"] = tau_IDvsElec_padded[:, 1]
events["tau1_id_vs_m"] = tau_IDvsMuon_padded[:, 0]
events["tau2_id_vs_m"] = tau_IDvsMuon_padded[:, 1]
events["tau1_id_vs_j"] = tau_IDvsJet_padded[:, 0]
events["tau2_id_vs_j"] = tau_IDvsJet_padded[:, 1]
return events
def _ak_to_numpy(ak_array, fields):
"""
Convert the given awkward array to a numpy table.
Parameters
----------
ak_array : awkward.Array
The awkward array, 2D or 3D.
fields : List
The column names of the last axis of the array.
Returns
-------
np_branch : tuple
Numpy-fied awkward array. See output of _branch_to_numpy.
"""
n_dims = ak_array.ndim - 1
if n_dims == 1:
n_items = np.ones(len(ak_array), dtype="int64")
elif n_dims == 2:
n_items = ak.num(ak_array).to_numpy()
ak_array = ak.flatten(ak_array)
else:
raise ValueError("Can not process array")
filled = np.ma.filled(
ak.pad_none(ak_array, target=len(fields), axis=-1).to_numpy(),
fill_value=np.nan,
)
return {fields[i]: filled[:, i] for i in range(len(fields))}, n_items
def association(cand1, cand2):
''' Function for association of the particles. The cuts that operates on all of them and
computation of quantities can go here. individual cuts can go on the main processing'''
#cut_dstar_back = ((cand2.deltamr > 0.143) & (cand2.deltamr < 0.148))
#cand2 = cand2[cut_dstar_back]
asso = ak.cartesian([cand1, cand2])
asso = asso[asso.slot0.vtxIdx == asso.slot1.vtxIdx]
asso = asso[ak.num(asso) > 0]
cand1 = ak.zip({
'pt': asso.slot0.pt,
'eta': asso.slot0.eta,
'phi': asso.slot0.phi,
'mass': asso.slot0.mass,
'charge': asso.slot0.charge}, with_name="PtEtaPhiMCandidate")
cand2 = ak.zip({
'pt': asso.slot1.pt,
'eta': asso.slot1.eta,
'phi': asso.slot1.phi,
'mass': asso.slot1.mass,
'charge': asso.slot1.charge}, with_name="PtEtaPhiMCandidate")
asso['deltarap'] = asso.slot0.rap - asso.slot1.rap
asso['cand'] = cand1 + cand2
return asso
def process(self, events):
# Note: This is required to ensure that behaviors are registered
# when running this code in a remote task.
ak.behavior.update(candidate.behavior)
output = self.accumulator.identity()
dataset = events.metadata['dataset']
muons = ak.zip(
{
"pt": events.Muon_pt,
"eta": events.Muon_eta,
"phi": events.Muon_phi,
"mass": events.Muon_mass,
"charge": events.Muon_charge,
},
with_name="PtEtaPhiMCandidate")
cut = (ak.num(muons) == 2) & (ak.sum(muons.charge) == 0)
# add first and second muon in every event together
dimuon = muons[cut][:, 0] + muons[cut][:, 1]
output["sumw"][dataset] += len(events)
output["mass"].fill(
dataset=dataset,
mass=dimuon.mass,
)
return output
def test_read_nanodata(suffix):
path = os.path.abspath(f'tests/samples/nano_dimuon.{suffix}')
factory =getattr(NanoEventsFactory, f'from_{suffix}')(path)
events = factory.events()
crossref(events)
crossref(events[ak.num(events.Jet) > 2])
def constituents(self, min_pt):
outputs_to_inputs = self.constituent_index(min_pt)
shape = ak.num(outputs_to_inputs)
total = np.sum(shape)
duplicate = ak.unflatten(np.zeros(total, np.int64), shape)
prepared = self.data[:, np.newaxis][duplicate]
return prepared[outputs_to_inputs]
def test_root_scalefactors():
extractor = lookup_tools.extractor()
extractor.add_weight_sets([
"testSF2d scalefactors_Tight_Electron tests/samples/testSF2d.histo.root"
])
extractor.finalize(reduce_list=["testSF2d"])
evaluator = extractor.make_evaluator()
counts, test_eta, test_pt = dummy_jagged_eta_pt()
# test flat eval
test_out = evaluator["testSF2d"](test_eta, test_pt)
# print it
print(evaluator["testSF2d"])
# test structured eval
test_eta_jagged = ak.unflatten(test_eta, counts)
test_pt_jagged = ak.unflatten(test_pt, counts)
test_out_jagged = evaluator["testSF2d"](test_eta_jagged, test_pt_jagged)
assert ak.all(ak.num(test_out_jagged) == counts)
assert ak.all(ak.flatten(test_out_jagged) == test_out)
print(test_out)
diff = np.abs(test_out - _testSF2d_expected_output)
print("Max diff: %.16f" % diff.max())
print("Median diff: %.16f" % np.median(diff))
print("Diff over threshold rate: %.1f %%" %
(100 * (diff >= 1.0e-8).sum() / diff.size))
assert (diff < 1.0e-8).all()
def test_read_nanomc(suffix):
path = os.path.abspath(f'tests/samples/nano_dy.{suffix}')
factory = getattr(NanoEventsFactory, f'from_{suffix}')(path)
events = factory.events()
# test after views first
genroundtrips(events.GenPart.mask[events.GenPart.eta > 0])
genroundtrips(events.mask[ak.any(events.Electron.pt > 50, axis=1)].GenPart)
genroundtrips(events.GenPart)
genroundtrips(events.GenPart[events.GenPart.eta > 0])
genroundtrips(events[ak.any(events.Electron.pt > 50, axis=1)].GenPart)
# sane gen matching (note for electrons gen match may be photon(22))
assert ak.all((abs(events.Electron.matched_gen.pdgId) == 11) | (events.Electron.matched_gen.pdgId == 22))
assert ak.all(abs(events.Muon.matched_gen.pdgId) == 13)
genroundtrips(events.Electron.matched_gen)
crossref(events[ak.num(events.Jet) > 2])
crossref(events)
if suffix == 'root':
assert ak.any(events.Photon.isTight, axis=1).tolist()[:9] == [False, True, True, True, False, False, False, False, False]
if suffix == 'parquet':
assert ak.any(events.Photon.isTight, axis=1).tolist()[:9] == [False, True, False, True, False, False, False, False, True]
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 compute_eff_algo(tau_1, tau_2, a, Pt_thr):
eff_presel = ak.sum(ak.num(tau_1, axis=-1))
deepTau_thr_1 = deep_thr(tau_1, a, Pt_thr)
deepTau_thr_2 = deep_thr(tau_2, a, Pt_thr)
deepTau_mask_1 = tau_1.deepTau_VSjet > deepTau_thr_1
deepTau_mask_2 = tau_2.deepTau_VSjet > deepTau_thr_2
eff = ak.sum(ditau_selection(deepTau_mask_1, deepTau_mask_2))
return compute_eff_witherr(eff, eff_presel)
def get_nnlo_weights(correction, events):
sl_evts = ak.num(events["SL"]) > 0
dl_evts = ak.num(events["DL"]) > 0
had_evts = ak.num(events["Had"]) > 0
# choose seed for random integers in similar fashio to https://github.com/CoffeaTeam/coffea/blob/master/coffea/jetmet_tools/CorrectedJetsFactory.py#L246-L254
seeds = np.array(ak.flatten(ak.concatenate([events["SL"][sl_evts]["TTbar"].pt, events["DL"][dl_evts]["TTbar"].pt, events["Had"][had_evts]["TTbar"].pt]), axis=None))[
[0, -1]
].view("i4")
randomstate = np.random.Generator(np.random.PCG64(seeds))
which_top_to_use = randomstate.integers(low=0, high=2, size=len(events)) # top is 0, tbar is 1
var = correction["Var"]
dist = correction["Correction"]
wts = np.ones(len(events))
if "thad_pt" in var:
# set wts for semilep evts
wts[sl_evts] = dist(ak.flatten(events["SL"][sl_evts]["THad"].pt, axis=None))
# set wts for dilep evts
dl_pt = np.where(which_top_to_use[dl_evts], ak.flatten(events["DL"][dl_evts]["Top"].pt, axis=None), ak.flatten(events["DL"][dl_evts]["Tbar"].pt, axis=None))
wts[dl_evts] = dist(dl_pt)
# set wts for had evts
had_pt = np.where(which_top_to_use[had_evts], ak.flatten(events["Had"][had_evts]["Top"].pt, axis=None), ak.flatten(events["Had"][had_evts]["Tbar"].pt, axis=None))
wts[had_evts] = dist(had_pt)
elif "mtt_vs_thad_ctstar" in var:
# set wts for semilep evts
thad_ctstar, tlep_ctstar = make_vars.ctstar(events["SL"][sl_evts]["THad"], events["SL"][sl_evts]["TLep"])
thad_ctstar, tlep_ctstar = ak.flatten(thad_ctstar, axis=None), ak.flatten(tlep_ctstar, axis=None)
wts[sl_evts] = dist(ak.flatten(events["SL"][sl_evts]["TTbar"].mass, axis=None), thad_ctstar)
# set wts for dilep evts
dl_top_ctstar, dl_tbar_ctstar = make_vars.ctstar(events["DL"][dl_evts]["Top"], events["DL"][dl_evts]["Tbar"])
dl_top_ctstar, dl_tbar_ctstar = ak.flatten(dl_top_ctstar, axis=None), ak.flatten(dl_tbar_ctstar, axis=None)
dl_ctstar = np.where(which_top_to_use[dl_evts], dl_top_ctstar, dl_tbar_ctstar)
wts[dl_evts] = dist(ak.flatten(events["DL"][dl_evts]["TTbar"].mass, axis=None), dl_ctstar)
# set wts for had evts
had_top_ctstar, had_tbar_ctstar = make_vars.ctstar(events["Had"][had_evts]["Top"], events["Had"][had_evts]["Tbar"])
had_top_ctstar, had_tbar_ctstar = ak.flatten(had_top_ctstar, axis=None), ak.flatten(had_tbar_ctstar, axis=None)
had_ctstar = np.where(which_top_to_use[had_evts], had_top_ctstar, had_tbar_ctstar)
wts[had_evts] = dist(ak.flatten(events["Had"][had_evts]["TTbar"].mass, axis=None), had_ctstar)
else:
raise ValueError("%s not supported for NNLO kinematic reweighting" % var)
return wts
def get_lepton_filter(events):
gp = events.GenPart
gp_lep = gp[((abs(gp.pdgId) == 11) | (abs(gp.pdgId) == 13) |
(abs(gp.pdgId) == 15))]
gp_lep_fromW = gp_lep[abs(gp_lep.parent.pdgId) == 24]
gp_leptonic_W_parent = find_first_parent(gp_lep_fromW.parent, maxgen=19)
gp_lep_fromW_noTop = gp_lep_fromW[abs(gp_leptonic_W_parent) != 6]
return ak.num(gp_lep_fromW_noTop) > 0
def test_num_9():
content = ak.layout.NumpyArray(
np.array([0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9]))
index = ak.layout.Index32(
np.array([0, 2, 4, 6, 8, 9, 7, 5], dtype=np.int64))
indexedarray = ak.Array(ak.layout.IndexedArray32(index, content))
cuda_indexedarray = ak.to_kernels(indexedarray, "cuda")
assert ak.num(cuda_indexedarray, 0) == ak.num(indexedarray, 0)
ioa = ak.Array(
ak.layout.IndexedOptionArray32(
ak.layout.Index32([-30, 19, 6, 7, -3, 21, 13, 22, 17, 9, -12, 16]),
ak.layout.NumpyArray(
np.array([
5.2,
1.7,
6.7,
-0.4,
4.0,
7.8,
3.8,
6.8,
4.2,
0.3,
4.6,
6.2,
6.9,
-0.7,
3.9,
1.6,
8.7,
-0.7,
3.2,
4.3,
4.0,
5.8,
4.2,
7.0,
5.6,
3.8,
])),
))
cuda_ioa = ak.to_kernels(ioa, "cuda")
ak.to_kernels(cuda_ioa, "cpu")
assert ak.num(cuda_ioa, 0) == ak.num(ioa, 0)
def add_vars(taus, c_type):
"""
Add variables to `taus` array for a given constituent type `c_type`
"""
taus[f'n_{c_type}'] = ak.num(
taus[c_type]) # counting number of constituents for each tau
for dim in ['phi', 'eta']:
taus[c_type, f'd{dim}'] = taus[c_type, dim] - taus[
f'tau_{dim}'] # normalising constituent coordinates wrt. tau direction
def test_read_nanomc(suffix):
path = os.path.abspath(f"tests/samples/nano_dy.{suffix}")
# parquet files were converted from even older nanoaod
nanoversion = NanoAODSchema.v6 if suffix == "root" else NanoAODSchema.v5
factory = getattr(NanoEventsFactory, f"from_{suffix}")(
path, schemaclass=nanoversion
)
events = factory.events()
# test after views first
genroundtrips(events.GenPart.mask[events.GenPart.eta > 0])
genroundtrips(events.mask[ak.any(events.Electron.pt > 50, axis=1)].GenPart)
genroundtrips(events.GenPart)
genroundtrips(events.GenPart[events.GenPart.eta > 0])
genroundtrips(events[ak.any(events.Electron.pt > 50, axis=1)].GenPart)
# sane gen matching (note for electrons gen match may be photon(22))
assert ak.all(
(abs(events.Electron.matched_gen.pdgId) == 11)
| (events.Electron.matched_gen.pdgId == 22)
)
assert ak.all(abs(events.Muon.matched_gen.pdgId) == 13)
genroundtrips(events.Electron.matched_gen)
crossref(events[ak.num(events.Jet) > 2])
crossref(events)
# test issue 409
assert ak.to_list(events[[]].Photon.mass) == []
if suffix == "root":
assert ak.any(events.Photon.isTight, axis=1).tolist()[:9] == [
False,
True,
True,
True,
False,
False,
False,
False,
False,
]
if suffix == "parquet":
assert ak.any(events.Photon.isTight, axis=1).tolist()[:9] == [
False,
True,
False,
True,
False,
False,
False,
False,
True,
]
def add_cuts(self, cuts, names):
for cut, name in zip(cuts, names):
if self.debug > 0:
n_objects_cut = awkward.sum(awkward.num(self.objects[cut]))
if self.n_objects_initial == 0:
return
print(
"%s ObjectCutDiagnostics: Applying cut %s would have an eff of %.4f"
% (self.cut_set, name,
float(n_objects_cut) / float(self.n_objects_initial)))
