def test_IndexedOptionArray_combinations():
content = awkward1.from_iter([[[0, 1, 2], [], [3, 4]], [], [[5]], [[6, 7, 8, 9]], [[], [10, 11, 12]]], highlevel=False)
index = awkward1.layout.Index64(numpy.array([0, 1, -1, -1, 4], dtype=numpy.int64))
array = awkward1.Array(awkward1.layout.IndexedOptionArray64(index, content))
assert awkward1.to_list(array) == [[[0, 1, 2], [], [3, 4]], [], None, None, [[], [10, 11, 12]]]
assert awkward1.to_list(awkward1.combinations(array, 2, axis=0)) == [([[0, 1, 2], [], [3, 4]], []), ([[0, 1, 2], [], [3, 4]], None), ([[0, 1, 2], [], [3, 4]], None), ([[0, 1, 2], [], [3, 4]], [[], [10, 11, 12]]), ([], None), ([], None), ([], [[], [10, 11, 12]]), (None, None), (None, [[], [10, 11, 12]]), (None, [[], [10, 11, 12]])]
assert awkward1.to_list(awkward1.combinations(array, 2, axis=1)) == [[([0, 1, 2], []), ([0, 1, 2], [3, 4]), ([], [3, 4])], [], None, None, [([], [10, 11, 12])]]
assert awkward1.to_list(awkward1.combinations(array, 2, axis=2)) == [[[(0, 1), (0, 2), (1, 2)], [], [(3, 4)]], [], None, None, [[], [(10, 11), (10, 12), (11, 12)]]]
def test_ByteMaskedArray_combinations():
content = awkward1.from_iter([[[0, 1, 2], [], [3, 4]], [], [[5]], [[6, 7, 8, 9]], [[], [10, 11, 12]]], highlevel=False)
mask = awkward1.layout.Index8(numpy.array([0, 0, 1, 1, 0], dtype=numpy.int8))
array = awkward1.Array(awkward1.layout.ByteMaskedArray(mask, content, valid_when=False))
assert awkward1.to_list(array) == [[[0, 1, 2], [], [3, 4]], [], None, None, [[], [10, 11, 12]]]
assert awkward1.to_list(awkward1.combinations(array, 2, axis=0)) == [([[0, 1, 2], [], [3, 4]], []), ([[0, 1, 2], [], [3, 4]], None), ([[0, 1, 2], [], [3, 4]], None), ([[0, 1, 2], [], [3, 4]], [[], [10, 11, 12]]), ([], None), ([], None), ([], [[], [10, 11, 12]]), (None, None), (None, [[], [10, 11, 12]]), (None, [[], [10, 11, 12]])]
assert awkward1.to_list(awkward1.combinations(array, 2, axis=1)) == [[([0, 1, 2], []), ([0, 1, 2], [3, 4]), ([], [3, 4])], [], None, None, [([], [10, 11, 12])]]
assert awkward1.to_list(awkward1.combinations(array, 2, axis=2)) == [[[(0, 1), (0, 2), (1, 2)], [], [(3, 4)]], [], None, None, [[], [(10, 11), (10, 12), (11, 12)]]]
def test_axis0():
array = awkward1.Array([0.0, 1.1, 2.2, 3.3])
assert awkward1.to_list(awkward1.combinations(array, 2, replacement=False, axis=0)) == [(0.0, 1.1), (0.0, 2.2), (0.0, 3.3), (1.1, 2.2), (1.1, 3.3), (2.2, 3.3)]
assert awkward1.to_list(awkward1.combinations(array, 2, replacement=False, axis=0, fields=["x", "y"])) == [{"x": 0.0, "y": 1.1}, {"x": 0.0, "y": 2.2}, {"x": 0.0, "y": 3.3}, {"x": 1.1, "y": 2.2}, {"x": 1.1, "y": 3.3}, {"x": 2.2, "y": 3.3}]
assert awkward1.combinations(array, 2, replacement=False, axis=0, parameters={"some": "param"}).layout.parameters["some"] == "param"
assert awkward1.to_list(awkward1.combinations(array, 3, replacement=False, axis=0)) == [(0.0, 1.1, 2.2), (0.0, 1.1, 3.3), (0.0, 2.2, 3.3), (1.1, 2.2, 3.3)]
def process(self, events):
# Initialize accumulator
out = self.accumulator.identity()
# Event selection: opposite charged same flavor
Electron = events.Electron
Electron_mask = (Electron.pt > 20) & (np.abs(Electron.eta) <
2.5) & (Electron.cutBased > 1)
Ele_channel_mask = ak.num(Electron[Electron_mask]) > 1
Ele_channel_events = events[Ele_channel_mask]
Ele = Ele_channel_events.Electron
# All possible pairs of Electron in each event
ele_pairs = ak.combinations(Ele, 2, axis=1)
# TLorentz vector sum of ele_pairs
ele_left, ele_right = ak.unzip(ele_pairs)
diele = ele_left + ele_right
diffsign_diele = diele[diele.charge == 0]
leading_diffsign_diele = diffsign_diele[ak.argmax(diffsign_diele.pt,
axis=1,
keepdims=True)]
#Mee = ak.flatten(leading_diffsign_diele.mass) # This makes type error ( primitive expected but ?float given )
Mee = ak.to_numpy(leading_diffsign_diele.mass)
Mee = Mee.flatten()
out.fill(dataset=events.metadata["dataset"], mass=Mee)
return out
def choose(first, n=2):
tmp = ak.combinations(first, n)
combs = tmp['0']
for i in range(1, n):
combs = combs.__add__(tmp[str(i)])
for i in range(n):
combs[str(i)] = tmp[str(i)]
return combs
def choose3(first, n=3):
from warnings import warn
warn(
"Deprecation Warning: The choose3 function will be removed. Use choose(first, n=3) instead."
)
tmp = ak.combinations(first, n)
combs = (tmp['0'] + tmp['1'] + tmp['2'])
combs['0'] = tmp['0']
combs['1'] = tmp['1']
combs['2'] = tmp['2']
return combs
def process(self, events):
output = self.accumulator.identity()
dataset = events.metadata["dataset"]
dimuon = ak.combinations(events.Muon, 2)
dimuon = dimuon["0"] + dimuon["1"]
output["pt"].fill(dataset=dataset, pt=ak.flatten(events.Muon.pt))
output["mass"].fill(dataset=dataset, mass=ak.flatten(dimuon.mass))
output["cutflow"]["%s_pt" % dataset] += sum(ak.num(events.Muon))
output["cutflow"]["%s_mass" % dataset] += sum(ak.num(dimuon))
return output
def test_ListOffsetArray():
array = awkward1.Array([[0.0, 1.1, 2.2, 3.3], [], [4.4, 5.5, 6.6], [7.7], [8.8, 9.9, 10.0, 11.1, 12.2]])
assert awkward1.to_list(awkward1.combinations(array, 2, replacement=False)) == [[(0.0, 1.1), (0.0, 2.2), (0.0, 3.3), (1.1, 2.2), (1.1, 3.3), (2.2, 3.3)], [], [(4.4, 5.5), (4.4, 6.6), (5.5, 6.6)], [], [(8.8, 9.9), (8.8, 10.0), (8.8, 11.1), (8.8, 12.2), (9.9, 10.0), (9.9, 11.1), (9.9, 12.2), (10.0, 11.1), (10.0, 12.2), (11.1, 12.2)]]
assert awkward1.to_list(awkward1.combinations(array, 2, replacement=False, fields=["x", "y"])) == [[{"x": 0.0, "y": 1.1}, {"x": 0.0, "y": 2.2}, {"x": 0.0, "y": 3.3}, {"x": 1.1, "y": 2.2}, {"x": 1.1, "y": 3.3}, {"x": 2.2, "y": 3.3}], [], [{"x": 4.4, "y": 5.5}, {"x": 4.4, "y": 6.6}, {"x": 5.5, "y": 6.6}], [], [{"x": 8.8, "y": 9.9}, {"x": 8.8, "y": 10.0}, {"x": 8.8, "y": 11.1}, {"x": 8.8, "y": 12.2}, {"x": 9.9, "y": 10.0}, {"x": 9.9, "y": 11.1}, {"x": 9.9, "y": 12.2}, {"x": 10.0, "y": 11.1}, {"x": 10.0, "y": 12.2}, {"x": 11.1, "y": 12.2}]]
tmp = awkward1.combinations(array, 2, replacement=False, parameters={"some": "param"}).layout
if isinstance(tmp, awkward1.partition.PartitionedArray):
assert awkward1.partition.first(tmp).content.parameters["some"] == "param"
else:
assert tmp.content.parameters["some"] == "param"
assert awkward1.to_list(awkward1.combinations(array, 2, replacement=True)) == [[(0.0, 0.0), (0.0, 1.1), (0.0, 2.2), (0.0, 3.3), (1.1, 1.1), (1.1, 2.2), (1.1, 3.3), (2.2, 2.2), (2.2, 3.3), (3.3, 3.3)], [], [(4.4, 4.4), (4.4, 5.5), (4.4, 6.6), (5.5, 5.5), (5.5, 6.6), (6.6, 6.6)], [(7.7, 7.7)], [(8.8, 8.8), (8.8, 9.9), (8.8, 10.0), (8.8, 11.1), (8.8, 12.2), (9.9, 9.9), (9.9, 10.0), (9.9, 11.1), (9.9, 12.2), (10.0, 10.0), (10.0, 11.1), (10.0, 12.2), (11.1, 11.1), (11.1, 12.2), (12.2, 12.2)]]
assert awkward1.to_list(awkward1.combinations(array, 3, replacement=False)) == [[(0.0, 1.1, 2.2), (0.0, 1.1, 3.3), (0.0, 2.2, 3.3), (1.1, 2.2, 3.3)], [], [(4.4, 5.5, 6.6)], [], [(8.8, 9.9, 10.0), (8.8, 9.9, 11.1), (8.8, 9.9, 12.2), (8.8, 10.0, 11.1), (8.8, 10.0, 12.2), (8.8, 11.1, 12.2), (9.9, 10.0, 11.1), (9.9, 10.0, 12.2), (9.9, 11.1, 12.2), (10.0, 11.1, 12.2)]]
assert awkward1.to_list(awkward1.combinations(array, 3, replacement=True)) == [[(0.0, 0.0, 0.0), (0.0, 0.0, 1.1), (0.0, 0.0, 2.2), (0.0, 0.0, 3.3), (0.0, 1.1, 1.1), (0.0, 1.1, 2.2), (0.0, 1.1, 3.3), (0.0, 2.2, 2.2), (0.0, 2.2, 3.3), (0.0, 3.3, 3.3), (1.1, 1.1, 1.1), (1.1, 1.1, 2.2), (1.1, 1.1, 3.3), (1.1, 2.2, 2.2), (1.1, 2.2, 3.3), (1.1, 3.3, 3.3), (2.2, 2.2, 2.2), (2.2, 2.2, 3.3), (2.2, 3.3, 3.3), (3.3, 3.3, 3.3)], [], [(4.4, 4.4, 4.4), (4.4, 4.4, 5.5), (4.4, 4.4, 6.6), (4.4, 5.5, 5.5), (4.4, 5.5, 6.6), (4.4, 6.6, 6.6), (5.5, 5.5, 5.5), (5.5, 5.5, 6.6), (5.5, 6.6, 6.6), (6.6, 6.6, 6.6)], [(7.7, 7.7, 7.7)], [(8.8, 8.8, 8.8), (8.8, 8.8, 9.9), (8.8, 8.8, 10.0), (8.8, 8.8, 11.1), (8.8, 8.8, 12.2), (8.8, 9.9, 9.9), (8.8, 9.9, 10.0), (8.8, 9.9, 11.1), (8.8, 9.9, 12.2), (8.8, 10.0, 10.0), (8.8, 10.0, 11.1), (8.8, 10.0, 12.2), (8.8, 11.1, 11.1), (8.8, 11.1, 12.2), (8.8, 12.2, 12.2), (9.9, 9.9, 9.9), (9.9, 9.9, 10.0), (9.9, 9.9, 11.1), (9.9, 9.9, 12.2), (9.9, 10.0, 10.0), (9.9, 10.0, 11.1), (9.9, 10.0, 12.2), (9.9, 11.1, 11.1), (9.9, 11.1, 12.2), (9.9, 12.2, 12.2), (10.0, 10.0, 10.0), (10.0, 10.0, 11.1), (10.0, 10.0, 12.2), (10.0, 11.1, 11.1), (10.0, 11.1, 12.2), (10.0, 12.2, 12.2), (11.1, 11.1, 11.1), (11.1, 11.1, 12.2), (11.1, 12.2, 12.2), (12.2, 12.2, 12.2)]]
def test_RegularArray():
array = awkward1.Array(numpy.array([[0.0, 1.1, 2.2, 3.3], [4.4, 5.5, 6.6, 7.7]]))
assert awkward1.to_list(awkward1.combinations(array, 2, replacement=False)) == [[(0.0, 1.1), (0.0, 2.2), (0.0, 3.3), (1.1, 2.2), (1.1, 3.3), (2.2, 3.3)], [(4.4, 5.5), (4.4, 6.6), (4.4, 7.7), (5.5, 6.6), (5.5, 7.7), (6.6, 7.7)]]
assert awkward1.to_list(awkward1.combinations(array, 2, replacement=False, fields=["x", "y"])) == [[{"x": 0.0, "y": 1.1}, {"x": 0.0, "y": 2.2}, {"x": 0.0, "y": 3.3}, {"x": 1.1, "y": 2.2}, {"x": 1.1, "y": 3.3}, {"x": 2.2, "y": 3.3}], [{"x": 4.4, "y": 5.5}, {"x": 4.4, "y": 6.6}, {"x": 4.4, "y": 7.7}, {"x": 5.5, "y": 6.6}, {"x": 5.5, "y": 7.7}, {"x": 6.6, "y": 7.7}]]
tmp = awkward1.combinations(array, 2, replacement=False, parameters={"some": "param"}).layout
if isinstance(tmp, awkward1.partition.PartitionedArray):
assert awkward1.partition.first(tmp).content.parameters["some"] == "param"
else:
assert tmp.content.parameters["some"] == "param"
assert awkward1.to_list(awkward1.combinations(array, 2, replacement=True)) == [[(0.0, 0.0), (0.0, 1.1), (0.0, 2.2), (0.0, 3.3), (1.1, 1.1), (1.1, 2.2), (1.1, 3.3), (2.2, 2.2), (2.2, 3.3), (3.3, 3.3)], [(4.4, 4.4), (4.4, 5.5), (4.4, 6.6), (4.4, 7.7), (5.5, 5.5), (5.5, 6.6), (5.5, 7.7), (6.6, 6.6), (6.6, 7.7), (7.7, 7.7)]]
assert awkward1.to_list(awkward1.combinations(array, 3, replacement=False)) == [[(0.0, 1.1, 2.2), (0.0, 1.1, 3.3), (0.0, 2.2, 3.3), (1.1, 2.2, 3.3)], [(4.4, 5.5, 6.6), (4.4, 5.5, 7.7), (4.4, 6.6, 7.7), (5.5, 6.6, 7.7)]]
assert awkward1.to_list(awkward1.combinations(array, 3, replacement=True)) == [[(0.0, 0.0, 0.0), (0.0, 0.0, 1.1), (0.0, 0.0, 2.2), (0.0, 0.0, 3.3), (0.0, 1.1, 1.1), (0.0, 1.1, 2.2), (0.0, 1.1, 3.3), (0.0, 2.2, 2.2), (0.0, 2.2, 3.3), (0.0, 3.3, 3.3), (1.1, 1.1, 1.1), (1.1, 1.1, 2.2), (1.1, 1.1, 3.3), (1.1, 2.2, 2.2), (1.1, 2.2, 3.3), (1.1, 3.3, 3.3), (2.2, 2.2, 2.2), (2.2, 2.2, 3.3), (2.2, 3.3, 3.3), (3.3, 3.3, 3.3)], [(4.4, 4.4, 4.4), (4.4, 4.4, 5.5), (4.4, 4.4, 6.6), (4.4, 4.4, 7.7), (4.4, 5.5, 5.5), (4.4, 5.5, 6.6), (4.4, 5.5, 7.7), (4.4, 6.6, 6.6), (4.4, 6.6, 7.7), (4.4, 7.7, 7.7), (5.5, 5.5, 5.5), (5.5, 5.5, 6.6), (5.5, 5.5, 7.7), (5.5, 6.6, 6.6), (5.5, 6.6, 7.7), (5.5, 7.7, 7.7), (6.6, 6.6, 6.6), (6.6, 6.6, 7.7), (6.6, 7.7, 7.7), (7.7, 7.7, 7.7)]]
final = final[(ak.count(qqb.pdgId, axis=1) == 3)]
finaljets = final.Jet
qqb = qqb[(ak.count(qqb.pdgId, axis=1) == 3)]
#Implementing Tight Jet Cuts on Training Data
finaljetSel = (abs(finaljets.eta) < 2.4) & (finaljets.pt > 30)
finalJets = finaljets[finaljetSel]
#Use nearest to match Jets
matchedGenJets = qqb.nearest(final.GenJet)
matchedJets = matchedGenJets.nearest(finalJets)
print("matched genpart to genjet and finally to reco jets")
#Assigning True false to sets of 3 jets
test = matchedJets.genJetIdx
combs = ak.combinations(finalJets, 3, replacement=False)
t1 = (combs['0'].genJetIdx == test[:, 0]) | (
combs['0'].genJetIdx == test[:, 1]) | (combs['0'].genJetIdx == test[:, 2])
t2 = (combs['1'].genJetIdx == test[:, 0]) | (
combs['1'].genJetIdx == test[:, 1]) | (combs['1'].genJetIdx == test[:, 2])
t3 = (combs['2'].genJetIdx == test[:, 0]) | (
combs['2'].genJetIdx == test[:, 1]) | (combs['2'].genJetIdx == test[:, 2])
t = t1 & t2 & t3
trutharray = ak.flatten(t)
print("matching a validation array for every combo of 3 jets")
#Zipping into CSV for training
jetcombos = ak.flatten(combs)
j1, j2, j3 = ak.unzip(jetcombos)
dR1_2 = j1.delta_r(j2)
def test_axis2():
array = awkward1.Array([[[0.0, 1.1, 2.2, 3.3], [], [4.4, 5.5, 6.6]], [], [[7.7], [8.8, 9.9, 10.0, 11.1, 12.2]]])
assert awkward1.to_list(awkward1.combinations(array, 2, axis=1, replacement=False)) == [[([0.0, 1.1, 2.2, 3.3], []), ([0.0, 1.1, 2.2, 3.3], [4.4, 5.5, 6.6]), ([], [4.4, 5.5, 6.6])], [], [([7.7], [8.8, 9.9, 10.0, 11.1, 12.2])]]
assert awkward1.to_list(awkward1.combinations(array, 2, axis=2, replacement=False)) == [[[(0.0, 1.1), (0.0, 2.2), (0.0, 3.3), (1.1, 2.2), (1.1, 3.3), (2.2, 3.3)], [], [(4.4, 5.5), (4.4, 6.6), (5.5, 6.6)]], [], [[], [(8.8, 9.9), (8.8, 10.0), (8.8, 11.1), (8.8, 12.2), (9.9, 10.0), (9.9, 11.1), (9.9, 12.2), (10.0, 11.1), (10.0, 12.2), (11.1, 12.2)]]]
def test_IndexedArray():
array = awkward1.Array([[0.0, 1.1, 2.2, 3.3], [], [4.4, 5.5, 6.6], None, [7.7], None, [8.8, 9.9, 10.0, 11.1, 12.2]])
assert awkward1.to_list(awkward1.combinations(array, 2, replacement=False)) == [[(0.0, 1.1), (0.0, 2.2), (0.0, 3.3), (1.1, 2.2), (1.1, 3.3), (2.2, 3.3)], [], [(4.4, 5.5), (4.4, 6.6), (5.5, 6.6)], None, [], None, [(8.8, 9.9), (8.8, 10.0), (8.8, 11.1), (8.8, 12.2), (9.9, 10.0), (9.9, 11.1), (9.9, 12.2), (10.0, 11.1), (10.0, 12.2), (11.1, 12.2)]]
def choose(first, n=2):
tmp = ak.combinations(first, n)
combs = (tmp['0'] + tmp['1'])
combs['0'] = tmp['0']
combs['1'] = tmp['1']
return combs
def process(self, events):
# Initialize accumulator
out = self.accumulator.identity()
dataset = setname
#events.metadata['dataset']
isData = 'genWeight' not in events.fields
selection = processor.PackedSelection()
# Cut flow
cut0 = np.zeros(len(events))
# --- Selection
# << flat dim helper function >>
def flat_dim(arr):
sub_arr = ak.flatten(arr)
mask = ~ak.is_none(sub_arr)
return ak.to_numpy(sub_arr[mask])
# << drop na helper function >>
def drop_na(arr):
mask = ~ak.is_none(arr)
return arr[mask]
# << drop na helper function >>
def drop_na_np(arr):
mask = ~np.isnan(arr)
return arr[mask]
# double lepton trigger
is_double_ele_trigger=True
if not is_double_ele_trigger:
double_ele_triggers_arr=np.ones(len(events), dtype=np.bool)
else:
double_ele_triggers_arr = np.zeros(len(events), dtype=np.bool)
for path in self._doubleelectron_triggers[self._year]:
if path not in events.HLT.fields: continue
double_ele_triggers_arr = double_ele_triggers_arr | events.HLT[path]
# single lepton trigger
is_single_ele_trigger=True
if not is_single_ele_trigger:
single_ele_triggers_arr=np.ones(len(events), dtype=np.bool)
else:
single_ele_triggers_arr = np.zeros(len(events), dtype=np.bool)
for path in self._singleelectron_triggers[self._year]:
if path not in events.HLT.fields: continue
single_ele_triggers_arr = single_ele_triggers_arr | events.HLT[path]
Initial_events = events
print("#### Initial events: ",Initial_events)
#events = events[single_ele_triggers_arr | double_ele_triggers_arr]
events = events[double_ele_triggers_arr]
##----------- Cut flow1: Passing Triggers
cut1 = np.ones(len(events))
print("#### cut1: ",len(cut1))
# Particle Identification
Electron = events.Electron
def Electron_selection(ele):
return(ele.pt > 25) & (np.abs(ele.eta) < 2.5) & (ele.cutBased > 2)
# Electron channel
Electron_mask = Electron_selection(Electron)
Ele_channel_mask = ak.num(Electron[Electron_mask]) > 1
Ele_channel_events = events[Ele_channel_mask]
##----------- Cut flow2: Electron channel
cut2 = np.ones(len(Ele_channel_events)) * 2
print("#### cut2: ",len(cut2))
# --- Calculate Scale factor weight
if not isData:
# PU weight with lookup table <-- On developing -->
#get_pu_weight = self._corrections['get_pu_weight'][self._year]
#pu = get_pu_weight(events.Pileup.nTrueInt)
get_ele_reco_sf = self._corrections['get_ele_reco_sf'][self._year]
get_ele_loose_id_sf = self._corrections['get_ele_loose_id_sf'][self._year]
get_ele_trig_leg1_SF = self._corrections['get_ele_trig_leg1_SF'][self._year]
get_ele_trig_leg1_data_Eff = self._corrections['get_ele_trig_leg1_data_Eff'][self._year]
get_ele_trig_leg1_mc_Eff = self._corrections['get_ele_trig_leg1_mc_Eff'][self._year]
get_ele_trig_leg2_SF = self._corrections['get_ele_trig_leg2_SF'][self._year]
get_ele_trig_leg2_data_Eff = self._corrections['get_ele_trig_leg2_data_Eff'][self._year]
get_ele_trig_leg2_mc_Eff = self._corrections['get_ele_trig_leg2_mc_Eff'][self._year]
# PU weight with custom made npy and multi-indexing
pu_weight_idx = ak.values_astype(Ele_channel_events.Pileup.nTrueInt,"int64")
pu = self._puweight_arr[pu_weight_idx]
nPV = Ele_channel_events.PV.npvsGood
else:
nPV = Ele_channel_events.PV.npvsGood
# Electron array
Ele = Ele_channel_events.Electron
Electron_mask = Electron_selection(Ele)
Ele_sel = Ele[Electron_mask]
# Electron pair
ele_pairs = ak.combinations(Ele_sel,2,axis=1)
ele_left, ele_right = ak.unzip(ele_pairs)
diele = ele_left + ele_right
# OS
os_mask = diele.charge == 0
os_diele = diele[os_mask]
os_ele_left = ele_left[os_mask]
os_ele_right = ele_right[os_mask]
os_event_mask = ak.num(os_diele) > 0
Ele_os_channel_events = Ele_channel_events[os_event_mask]
#selection.add('ossf',os_event_mask)
# Helper function: High PT argmax
def make_leading_pair(target,base):
return target[ak.argmax(base.pt,axis=1,keepdims=True)]
# -- Only Leading pair --
leading_diele = make_leading_pair(diele,diele)
leading_ele = make_leading_pair(ele_left,diele)
subleading_ele= make_leading_pair(ele_right,diele)
# -- Scale Factor for each electron
def Trigger_Weight(eta1,pt1,eta2,pt2):
per_ev_MC =\
get_ele_trig_leg1_mc_Eff(eta1,pt1) * get_ele_trig_leg2_mc_Eff(eta2,pt2) +\
get_ele_trig_leg1_mc_Eff(eta2,pt2) * get_ele_trig_leg2_mc_Eff(eta1,pt1) -\
get_ele_trig_leg1_mc_Eff(eta1,pt1) * get_ele_trig_leg1_mc_Eff(eta2,pt2)
per_ev_data =\
get_ele_trig_leg1_data_Eff(eta1,pt1) * get_ele_trig_leg1_SF(eta1,pt1) * get_ele_trig_leg2_data_Eff(eta2,pt2) * get_ele_trig_leg2_SF(eta2,pt2) +\
get_ele_trig_leg1_data_Eff(eta2,pt2) * get_ele_trig_leg1_SF(eta2,pt2) * get_ele_trig_leg2_data_Eff(eta1,pt1) * get_ele_trig_leg2_SF(eta1,pt1) -\
get_ele_trig_leg1_data_Eff(eta1,pt1) * get_ele_trig_leg1_SF(eta1,pt1) * get_ele_trig_leg1_data_Eff(eta2,pt2) * get_ele_trig_leg1_SF(eta2,pt2)
return per_ev_data/per_ev_MC
if not isData:
ele_loose_id_sf = get_ele_loose_id_sf(ak.flatten(leading_ele.deltaEtaSC + leading_ele.eta),ak.flatten(leading_ele.pt))* get_ele_loose_id_sf(ak.flatten(subleading_ele.deltaEtaSC + subleading_ele.eta),ak.flatten(subleading_ele.pt))
#print("Ele ID SC---->",ele_loose_id_sf)
ele_reco_sf = get_ele_reco_sf(ak.flatten(leading_ele.deltaEtaSC + leading_ele.eta),ak.flatten(leading_ele.pt))* get_ele_reco_sf(ak.flatten(subleading_ele.deltaEtaSC + subleading_ele.eta),ak.flatten(subleading_ele.pt))
#print("Ele RECO SC---->",ele_reco_sf)
eta1 = ak.flatten(leading_ele.deltaEtaSC + leading_ele.eta)
eta2 = ak.flatten(subleading_ele.deltaEtaSC + subleading_ele.eta)
pt1 = ak.flatten(leading_ele.pt)
pt2 = ak.flatten(subleading_ele.pt)
ele_trig_weight = Trigger_Weight(eta1,pt1,eta2,pt2)
print("#### Test print trigger weight ####")
print(ele_trig_weight)
# --OS and Leading pair --
leading_os_diele = make_leading_pair(os_diele,os_diele)
leading_os_ele = make_leading_pair(os_ele_left,os_diele)
subleading_os_ele= make_leading_pair(os_ele_right,os_diele)
##----------- Cut flow3: OSSF
cut3 = np.ones(len(flat_dim(leading_os_diele))) * 3
print("#### cut3: ",len(cut3))
# Helper function: Zmass window
def makeZmass_window_mask(dielecs,start=60,end=120):
mask = (dielecs.mass >= start) & (dielecs.mass <= end)
return mask
# -- OS and Leading pair --
Zmass_mask_os = makeZmass_window_mask(leading_os_diele)
leading_os_Zwindow_ele = leading_os_ele[Zmass_mask_os]
subleading_os_Zwindow_ele = subleading_os_ele[Zmass_mask_os]
leading_os_Zwindow_diele = leading_os_diele[Zmass_mask_os]
# for masking
Zmass_event_mask = makeZmass_window_mask(leading_diele)
Zmass_os_event_mask= ak.flatten(os_event_mask * Zmass_event_mask)
Ele_Zmass_os_events = Ele_channel_events[Zmass_os_event_mask]
##----------- Cut flow4: Zmass
cut4 = np.ones(len(flat_dim(leading_os_Zwindow_diele))) * 4
print("#### cut4: ",len(cut4))
## << Selection method -- Need validation >>
#print("a--->",len(Ele_channel_events))
#print("b--->",len(Ele_os_channel_events))
#print("b2--->",len(cut3))
#print("c--->",len(Ele_Zmass_os_events))
#print("c2--->",len(cut4))
ele1PT = flat_dim(leading_os_Zwindow_ele.pt)
ele1Eta = flat_dim(leading_os_Zwindow_ele.eta)
ele1Phi = flat_dim(leading_os_Zwindow_ele.phi)
ele2PT = flat_dim(subleading_os_Zwindow_ele.pt)
ele2Eta = flat_dim(subleading_os_Zwindow_ele.eta)
ele2Phi = flat_dim(subleading_os_Zwindow_ele.phi)
Mee = flat_dim(leading_os_Zwindow_diele.mass)
charge = flat_dim(leading_os_Zwindow_diele.charge)
# --- Apply weight and hist
weights = processor.Weights(len(cut2))
# --- skim cut-weight
def skim_weight(arr):
mask1 = ~ak.is_none(arr)
subarr = arr[mask1]
mask2 = subarr !=0
return ak.to_numpy(subarr[mask2])
cuts = ak.flatten(Zmass_mask_os)
if not isData:
weights.add('pileup',pu)
weights.add('ele_id',ele_loose_id_sf)
weights.add('ele_reco',ele_reco_sf)
#weights.add('ele_trigger',ele_trig_weight)
# Initial events
out["sumw"][dataset] += len(Initial_events)
# Cut flow loop
for cut in [cut0,cut1,cut2,cut3,cut4]:
out["cutflow"].fill(
dataset = dataset,
cutflow=cut
)
# Primary vertex
out['nPV'].fill(
dataset=dataset,
nPV = nPV,
weight = weights.weight()
)
out['nPV_nw'].fill(
dataset=dataset,
nPV_nw = nPV
)
# Physics varibles passing Zwindow
out["mass"].fill(
dataset=dataset,
mass=Mee,
weight = skim_weight(weights.weight() * cuts)
)
out["ele1pt"].fill(
dataset=dataset,
ele1pt=ele1PT,
weight = skim_weight(weights.weight() * cuts)
)
out["ele1eta"].fill(
dataset=dataset,
ele1eta=ele1Eta,
weight = skim_weight(weights.weight() * cuts)
)
out["ele1phi"].fill(
dataset=dataset,
ele1phi=ele1Phi,
weight = skim_weight(weights.weight() * cuts)
)
out["ele2pt"].fill(
dataset=dataset,
ele2pt=ele2PT,
weight = skim_weight(weights.weight() * cuts)
)
out["ele2eta"].fill(
dataset=dataset,
ele2eta=ele2Eta,
weight = skim_weight(weights.weight() * cuts)
)
out["ele2phi"].fill(
dataset=dataset,
ele2phi=ele2Phi,
weight = skim_weight(weights.weight() * cuts)
)
return out
#Open ROOT file from Pythia Delphes production with FCCSW
file = uproot.open(f"{path}/output/rootfiles/FCCDelphesOutput_100events.root")
tree = file['events'] #Get TTree from file
#Get the charged tracks
tr = tree.arrays(filter_name="pfcharged.core*")
tr["p"] = np.sqrt(tr["pfcharged.core.p4.px"]**2 + tr["pfcharged.core.p4.py"]**2 + tr["pfcharged.core.p4.pz"]**2)
tr["pt"] = np.sqrt(tr["pfcharged.core.p4.px"]**2 + tr["pfcharged.core.p4.py"]**2)
tr['eta'] = np.log((tr['p'] + tr['pfcharged.core.p4.pz'])/(tr['p'] - tr['pfcharged.core.p4.pz']))/2
tr['phi'] = np.arctan2(tr['pfcharged.core.p4.py'], tr['pfcharged.core.p4.px'])
#Pions
pi_cut = abs(tr["pfcharged.core.pdgId"]) == 211
pi = tr[pi_cut]
#Number of pions in each event
pi_sum = ak.num(pi_cut)
#Keep events with 2 or more pions
pi = pi[pi_sum >= 2]
#Make pion pairs per event
pi_pairs = ak.combinations(pi, 2)
#pt of first pion pair from first event
#print(pi_pairs[0]["pt"][0])
pi1, pi2 = ak.unzip(pi_pairs)
m_pipi = calc_invariant_mass(pi1, pi2)
plt.hist(ak.flatten(m_pipi),bins=400)
plt.show()
def process(self, events):
# Initialize accumulator
out = self.accumulator.identity()
dataset = setname
#events.metadata['dataset']
isData = 'genWeight' not in events.fields
selection = processor.PackedSelection()
# --- Calculate Scale factor weight
if not isData:
# PU weight
get_pu_weight = self._corrections['get_pu_weight'][self._year]
pu = get_pu_weight(events.Pileup.nTrueInt)
# --- Selection
# flat dim helper function
def flat_dim(arr):
sub_arr = ak.flatten(arr)
mask = ~ak.is_none(sub_arr)
return ak.to_numpy(sub_arr[mask])
# double lepton trigger
is_double_ele_trigger = True
if not is_double_ele_trigger:
double_ele_triggers_arr = np.ones(len(events), dtype=np.bool)
else:
double_ele_triggers_arr = np.zeros(len(events), dtype=np.bool)
for path in self._doubleelectron_triggers[self._year]:
if path not in events.HLT.fields: continue
double_ele_triggers_arr = double_ele_triggers_arr | events.HLT[
path]
#print("Double Electron triggers")
#print(double_ele_triggers_arr,len(double_ele_triggers_arr))
# single lepton trigger
is_single_ele_trigger = True
if not is_single_ele_trigger:
single_ele_triggers_arr = np.ones(len(events), dtype=np.bool)
else:
single_ele_triggers_arr = np.zeros(len(events), dtype=np.bool)
for path in self._singleelectron_triggers[self._year]:
if path not in events.HLT.fields: continue
single_ele_triggers_arr = single_ele_triggers_arr | events.HLT[
path]
#print("Single Electron triggers")
#print(single_ele_triggers_arr,len(single_ele_triggers_arr))
Initial_events = events
print("{0} number of events are detected".format(len(Initial_events)))
events = events[single_ele_triggers_arr | double_ele_triggers_arr]
print(
"Total {0} number of events are remain after triggger | Eff: {1}".
format(len(events),
len(events) / len(Initial_events) * 100))
# Particle Identification
Electron = events.Electron
def Electron_selection(ele):
return (ele.pt > 20) & (np.abs(ele.eta) < 2.5) & (ele.cutBased > 1)
# Electron channel
Electron_mask = Electron_selection(Electron)
Ele_channel_mask = ak.num(Electron[Electron_mask]) > 1
Ele_channel_events = events[Ele_channel_mask]
N_Ele_Channel_events = len(Ele_channel_events)
print("#1 Ele channel evts: {0} --> {1}".format(
len(events), N_Ele_Channel_events))
# Electron array
Ele = Ele_channel_events.Electron
Electron_mask = Electron_selection(Ele)
Ele_sel = Ele[Electron_mask]
# Electron pair
ele_pairs = ak.combinations(Ele_sel, 2, axis=1)
ele_left, ele_right = ak.unzip(ele_pairs)
diele = ele_left + ele_right
# OS
os_mask = diele.charge == 0
os_diele = diele[os_mask]
os_ele_left = ele_left[os_mask]
os_ele_right = ele_right[os_mask]
# Helper function: High PT argmax
def make_leading_pair(target, base):
return target[ak.argmax(base.pt, axis=1, keepdims=True)]
# -- Only Leading pair --
#leading_diele = make_leading_pair(diele,diele)
#leading_ele = make_leading_pair(ele_left,diele)
#subleading_ele= make_leading_pair(ele_right,diele)
# --OS and Leading pair --
leading_os_diele = make_leading_pair(os_diele, os_diele)
leading_os_ele = make_leading_pair(os_ele_left, os_diele)
subleading_os_ele = make_leading_pair(os_ele_right, os_diele)
# Helper function: Zmass window
def makeZmass_window_mask(dielecs, start=60, end=120):
mask = (dielecs.mass >= start) & (dielecs.mass <= end)
return mask
# -- Only Leading pair --
#Zmass_mask = makeZmass_window_mask(leading_diele)
#leading_Zwindow_ele = leading_ele[Zmass_mask]
#subleading_Zwindow_ele = subleading_ele[Zmass_mask]
#leading_Zwindow_diele = leading_diele[Zmass_mask]
# -- OS and Leading pair --
Zmass_mask_os = makeZmass_window_mask(leading_os_diele)
leading_os_Zwindow_ele = leading_os_ele[Zmass_mask_os]
subleading_os_Zwindow_ele = subleading_os_ele[Zmass_mask_os]
leading_os_Zwindow_diele = leading_os_diele[Zmass_mask_os]
#ele1PT = flat_dim(leading_Zwindow_ele.pt)
#ele1Eta = flat_dim(leading_Zwindow_ele.eta)
#ele1Phi = flat_dim(leading_Zwindow_ele.phi)
#ele2PT = flat_dim(subleading_Zwindow_ele.pt)
#ele2Eta = flat_dim(subleading_Zwindow_ele.eta)
#ele2Phi = flat_dim(subleading_Zwindow_ele.phi)
#Mee = flat_dim(leading_Zwindow_diele.mass)
#charge = flat_dim(leading_Zwindow_diele.charge)
os_ele1PT = flat_dim(leading_os_Zwindow_ele.pt)
os_ele1Eta = flat_dim(leading_os_Zwindow_ele.eta)
os_ele1Phi = flat_dim(leading_os_Zwindow_ele.phi)
os_ele2PT = flat_dim(subleading_os_Zwindow_ele.pt)
os_ele2Eta = flat_dim(subleading_os_Zwindow_ele.eta)
os_ele2Phi = flat_dim(subleading_os_Zwindow_ele.phi)
os_Mee = flat_dim(leading_os_Zwindow_diele.mass)
os_charge = flat_dim(leading_os_Zwindow_diele.charge)
#print("#5 Leading PT : {0}".format(len(ele1PT)))
#print("#5 Leading os PT {0}".format(len(os_ele1PT)))
# --- Apply weight --- on progress #
#if not isData:
# weights = processor.Weights(len(events))
# weights.add('pileup',pu)
out["sumw"][dataset] += len(events)
out["os_mass"].fill(dataset=dataset, os_mass=os_Mee)
out["os_ele1pt"].fill(dataset=dataset, os_ele1pt=os_ele1PT)
out["os_ele1eta"].fill(dataset=dataset, os_ele1eta=os_ele1Eta)
out["os_ele1phi"].fill(dataset=dataset, os_ele1phi=os_ele1Phi)
out["os_ele2pt"].fill(dataset=dataset, os_ele2pt=os_ele2PT)
out["os_ele2eta"].fill(dataset=dataset, os_ele2eta=os_ele2Eta)
out["os_ele2phi"].fill(dataset=dataset, os_ele2phi=os_ele2Phi)
return out
