def test_broadcast_and_apply_levels():
arrays = [
ak.Array([[[0.0, 1.1, 2.2], []], [[3.3, 4.4]],
[[5.5], [6.6, 7.7, 8.8, 9.9]]]),
ak.Array([[[10, 20], [30]], [[40]], [[50, 60, 70], [80, 90]]]),
]
# nothing is required to have the same length
assert ak.concatenate(arrays, axis=0).tolist() == [
[[0.0, 1.1, 2.2], []],
[[3.3, 4.4]],
[[5.5], [6.6, 7.7, 8.8, 9.9]],
[[10, 20], [30]],
[[40]],
[[50, 60, 70], [80, 90]],
]
# the outermost arrays are required to have the same length, but nothing deeper than that
assert ak.concatenate(arrays, axis=1).tolist() == [
[[0.0, 1.1, 2.2], [], [10, 20], [30]],
[[3.3, 4.4], [40]],
[[5.5], [6.6, 7.7, 8.8, 9.9], [50, 60, 70], [80, 90]],
]
# the outermost arrays and the first level are required to have the same length, but nothing deeper
assert ak.concatenate(arrays, axis=2).tolist() == [
[[0.0, 1.1, 2.2, 10, 20], [30]],
[[3.3, 4.4, 40]],
[[5.5, 50, 60, 70], [6.6, 7.7, 8.8, 9.9, 80, 90]],
]
def test_concatenate():
one = ak.Array([1.1, 2.2, 3.3, 4.4, 5.5], check_valid=True)
two = ak.Array([[], [1], [2, 2], [3, 3, 3]], check_valid=True)
three = ak.Array([True, False, False, True, True], check_valid=True)
assert ak.to_list(ak.concatenate([one, two, three])) == [
1.1,
2.2,
3.3,
4.4,
5.5,
[],
[1],
[2, 2],
[3, 3, 3],
1.0,
0.0,
0.0,
1.0,
1.0,
]
assert isinstance(ak.concatenate([one, two, three], highlevel=False),
ak.layout.UnionArray8_64)
assert len(ak.concatenate([one, two, three],
highlevel=False).contents) == 2
def test_listoffsetarray_concatenate():
content_one = ak.layout.NumpyArray(np.array([1, 2, 3, 4, 5, 6, 7, 8, 9]))
offsets_one = ak.layout.Index64(np.array([0, 3, 3, 5, 9]))
one = ak.layout.ListOffsetArray64(offsets_one, content_one)
assert ak.to_list(one) == [[1, 2, 3], [], [4, 5], [6, 7, 8, 9]]
content_two = ak.layout.NumpyArray(np.array([100, 200, 300, 400, 500]))
offsets_two = ak.layout.Index64(np.array([0, 2, 4, 4, 5]))
two = ak.layout.ListOffsetArray64(offsets_two, content_two)
assert ak.to_list(two) == [[100, 200], [300, 400], [], [500]]
assert ak.to_list(ak.concatenate([one, two], 0)) == [
[1, 2, 3],
[],
[4, 5],
[6, 7, 8, 9],
[100, 200],
[300, 400],
[],
[500],
]
assert ak.to_list(ak.concatenate([one, two], 1)) == [
[1, 2, 3, 100, 200],
[300, 400],
[4, 5],
[6, 7, 8, 9, 500],
]
def test_indexed_array_concatenate():
one = ak.Array([[1, 2, 3], [None, 4], None, [None, 5]]).layout
two = ak.Array([6, 7, 8]).layout
three = ak.Array([[6.6], [7.7, 8.8]]).layout
four = ak.Array([[6.6], [7.7, 8.8], None, [9.9]]).layout
assert ak.to_list(ak.concatenate([one, two], 0)) == [
[1, 2, 3],
[None, 4],
None,
[None, 5],
6,
7,
8,
]
with pytest.raises(ValueError):
ak.to_list(ak.concatenate([one, three], 1))
assert ak.to_list(ak.concatenate([one, four], 1)) == [
[1, 2, 3, 6.6],
[None, 4, 7.7, 8.8],
[],
[None, 5, 9.9],
]
def test_concatenate_number():
assert ak.to_list(
ak.concatenate([ak.Array([[1, 2, 3], [], [4, 5]]), 999],
axis=1)) == [[1, 2, 3, 999], [999], [4, 5, 999]]
assert ak.to_list(
ak.concatenate(
[ak.Array([[[1.1], [2.2, 3.3]], [[]], [[4.4], [5.5]]]), 999],
axis=2)) == [[[1.1, 999.0], [2.2, 3.3, 999.0]], [[999.0]],
[[4.4, 999.0], [5.5, 999.0]]]
assert (str(
ak.type(
ak.concatenate(
[
ak.Array([[1, 2, 3], [], [4, 5]]),
ak.Array([[123], [223], [323]]),
],
axis=1,
))) == "3 * var * int64")
assert ak.to_list(
ak.concatenate([
ak.Array([[1, 2, 3], [], [4, 5]]),
ak.Array([[123], [223], [323]])
],
axis=1)) == [[1, 2, 3, 123], [223], [4, 5, 323]]
one = ak.Array([[1, 2, 3], [], [4, 5]])
two = ak.Array([[123], [223], [323]])
assert ak.to_list(ak.concatenate([one, two], axis=1)) == [
[1, 2, 3, 123],
[223],
[4, 5, 323],
]
def test_negative_axis_concatenate():
arrays = [
ak.Array([[[0.0, 1.1, 2.2], []], [[3.3, 4.4]],
[[5.5], [6.6, 7.7, 8.8, 9.9]]]),
ak.Array([[[10, 20], [30]], [[40]], [[50, 60, 70], [80, 90]]]),
]
assert ak.concatenate(arrays, axis=-1).tolist() == [
[[0.0, 1.1, 2.2, 10, 20], [30]],
[[3.3, 4.4, 40]],
[[5.5, 50, 60, 70], [6.6, 7.7, 8.8, 9.9, 80, 90]],
]
assert ak.concatenate(arrays, axis=-2).tolist() == [
[[0.0, 1.1, 2.2], [], [10, 20], [30]],
[[3.3, 4.4], [40]],
[[5.5], [6.6, 7.7, 8.8, 9.9], [50, 60, 70], [80, 90]],
]
assert ak.concatenate(arrays, axis=-3).tolist() == [
[[0.0, 1.1, 2.2], []],
[[3.3, 4.4]],
[[5.5], [6.6, 7.7, 8.8, 9.9]],
[[10, 20], [30]],
[[40]],
[[50, 60, 70], [80, 90]],
]
def test_feature():
one = ak.Array([[{"x": 1}], [], [{"x": 2}]], with_name="One")
two = ak.Array([[{"x": 1.1}], [], [{"x": 2.2}]], with_name="Two")
assert (str(ak.with_name(ak.concatenate([one, two], axis=1),
"All").type) == '3 * var * All["x": float64]')
assert (str(
ak.with_name(ak.concatenate([one[1:], two[1:]], axis=1),
"All").type) == '2 * var * All["x": float64]')
def test_concatenate():
one = ak.Array(ak.layout.RecordArray([], [], length=0))
two = ak.Array(ak.layout.RecordArray([], [], length=0))
assert len(ak.concatenate([one, two])) == 0
assert len(ak.concatenate([one])) == 0
one = ak.Array(ak.layout.RecordArray([], [], length=3))
two = ak.Array(ak.layout.RecordArray([], [], length=5))
assert len(ak.concatenate([one, two])) == 8
assert len(ak.concatenate([one])) == 3
def test_jet_resolution():
from coffea.jetmet_tools import JetResolution
counts, test_eta, test_pt = dummy_jagged_eta_pt()
test_Rho = np.full_like(test_eta, 10.0)
test_pt_jag = ak.unflatten(test_pt, counts)
test_eta_jag = ak.unflatten(test_eta, counts)
test_Rho_jag = ak.unflatten(test_Rho, counts)
jer_names = ["Spring16_25nsV10_MC_PtResolution_AK4PFPuppi"]
reso = JetResolution(**{name: evaluator[name] for name in jer_names})
print(reso)
resos = reso.getResolution(JetEta=test_eta, Rho=test_Rho, JetPt=test_pt)
resos_jag = reso.getResolution(JetEta=test_eta_jag,
Rho=test_Rho_jag,
JetPt=test_pt_jag)
assert ak.all(np.abs(resos - ak.flatten(resos_jag)) < 1e-6)
test_pt_jag = test_pt_jag[0:3]
test_eta_jag = test_eta_jag[0:3]
test_Rho_jag = test_Rho_jag[0:3]
test_Rho_jag = ak.concatenate(
[test_Rho_jag[:-1], [ak.concatenate([test_Rho_jag[-1, :-1], 100.0])]])
counts = counts[0:3]
print("Raw jet values:")
print("pT:", test_pt_jag)
print("eta:", test_eta_jag)
print("rho:", test_Rho_jag, "\n")
resos_jag_ref = ak.unflatten(
np.array([
0.21974642,
0.32421591,
0.33702479,
0.27420327,
0.13940689,
0.48134521,
0.26564994,
1.0,
]),
counts,
)
resos_jag = reso.getResolution(JetEta=test_eta_jag,
Rho=test_Rho_jag,
JetPt=test_pt_jag)
print("Reference Resolution (jagged):", resos_jag_ref)
print("Resolution (jagged):", resos_jag)
# NB: 5e-4 tolerance was agreed upon by lgray and aperloff, if the differences get bigger over time
# we need to agree upon how these numbers are evaluated (double/float conversion is kinda random)
assert ak.all(
np.abs(ak.flatten(resos_jag_ref) - ak.flatten(resos_jag)) < 5e-4)
def test_numpyarray_concatenate_axis0():
np1 = np.arange(2 * 7 * 5, dtype=np.float64).reshape(2, 7, 5)
np2 = np.arange(3 * 7 * 5, dtype=np.int64).reshape(3, 7, 5)
ak1 = ak.layout.NumpyArray(np1)
ak2 = ak.layout.NumpyArray(np2)
assert ak.to_list(ak.concatenate([np1, np2, np1, np2], 0)) == ak.to_list(
np.concatenate([np1, np2, np1, np2], 0))
assert ak.to_list(ak.concatenate([ak1, ak2], 0)) == ak.to_list(
np.concatenate([ak1, ak2], 0))
assert ak.to_list(np.concatenate([ak1, ak2], 0)) == ak.to_list(
ak.concatenate([np1, np2], 0))
def append(self, rhs):
'''
like concatenate, axis=1
so that the track collection can be appended to the rechit collection
'''
self.splitIdx= ak1.concatenate([self.splitIdx,rhs.splitIdx],axis=1)
self._checkshapes(self.features,rhs.features)
self.features = ak1.concatenate([self.features, rhs.features],axis=1)
newtruth={}
for k in self.truth.keys():
self._checkshapes(self.truth[k],rhs.truth[k])
newtruth[k] = ak1.concatenate([self.truth[k], rhs.truth[k]],axis=1)
self.truth = newtruth
def to_awkward(self, **kwargs):
"""
See ``DatasetProvider.to_awkward()``. ``X`` and ``y`` will
be concatenated along axis=0 (one provider after another)
"""
all_X = []
all_y = []
for p in self.providers:
X, y = p.to_awkward(**kwargs)
all_X.append(X)
all_y.append(y)
return ak.concatenate(all_X), ak.concatenate(all_y)
def test():
a = ak.Array([{"this": 100}])
b = ak.Array([{"this": 90, "that": 100}])
c = ak.concatenate((a, b))
with pytest.raises(ValueError):
ak.unzip(c)
a = ak.Array([{"this": 100}])
b = ak.Array([{"this": 90}])
c = ak.concatenate((a, b))
(tmp, ) = ak.unzip(c)
assert tmp.tolist() == [100, 90]
def _append_object(event_list, field):
new_event_list = []
for i in range(len(event_list)):
event_list_i = awkward.fromiter(event_list[i])
field_i = awkward.fromiter(field[i])
new_event_list.append(awkward.concatenate([event_list_i, field_i], axis=1).tolist())
return new_event_list
def combine_awkward(self, awks: Iterable[ak.Array]) -> ak.Array:
'''Combine many awkward arrays into a single one, in order.
Args:
awks (Iterable[ChunkedArray]): The input list of awkward arrays
'''
return ak.concatenate(awks) # type: ignore
def test_cartesian():
muon = ak.Array([[{"pt": 1.0}], []], with_name="muon")
electron = ak.Array([[], [{"pt": 1.0}]], with_name="electron")
muon = muon[muon.pt > 5]
electron = electron[electron.pt > 5]
leptons = ak.concatenate([muon, electron], axis=1)
candidate = ak.firsts(leptons)
assert ak.to_list(ak.Array(candidate)) == [None, None]
result = ak.cartesian([candidate, candidate], axis=0)
assert ak.to_list(result) == [
(None, None),
(None, None),
(None, None),
(None, None),
]
result = ak.cartesian([candidate, ak.Array([[1, 2, 3], []])], axis=1)
assert ak.to_list(result) == [None, None]
one, two = ak.broadcast_arrays(candidate, ak.Array([[1, 2, 3], []]))
assert ak.to_list(one) == [None, None]
assert ak.to_list(two) == [None, None]
def _splitJaggedArray(self, jagged):
if self.splitIdx is None:
raise ValueError("First determine split indices by running _readSplits")
split1 = jagged[self.splitIdx]
split2 = jagged[~self.splitIdx]
arr = ak1.concatenate([split1,split2],axis=0)
return arr
def process(self, events):
events["Electron", "pdgId"] = -11 * events.Electron.charge
events["Muon", "pdgId"] = -13 * events.Muon.charge
events["leptons"] = ak.concatenate(
[events.Electron, events.Muon],
axis=1,
)
events = events[ak.num(events.leptons) >= 3]
pair = ak.argcombinations(events.leptons, 2, fields=["l1", "l2"])
pair = pair[(
events.leptons[pair.l1].pdgId == -events.leptons[pair.l2].pdgId)]
with np.errstate(invalid="ignore"):
pair = pair[ak.singletons(
ak.argmin(
abs((events.leptons[pair.l1] +
events.leptons[pair.l2]).mass - 91.2),
axis=1,
))]
events = events[ak.num(pair) > 0]
pair = pair[ak.num(pair) > 0][:, 0]
l3 = ak.local_index(events.leptons)
l3 = l3[(l3 != pair.l1) & (l3 != pair.l2)]
l3 = l3[ak.argmax(events.leptons[l3].pt, axis=1, keepdims=True)]
l3 = events.leptons[l3][:, 0]
mt = np.sqrt(2 * l3.pt * events.MET.pt *
(1 - np.cos(events.MET.delta_phi(l3))))
return (hist.Hist.new.Reg(
100, 0, 200, name="mt",
label=r"$\ell$-MET transverse mass [GeV]").Double().fill(mt))
def _concat(arrays, axis=0):
if len(arrays) == 0:
return np.array([])
if isinstance(arrays[0], np.ndarray):
return np.concatenate(arrays, axis=axis)
else:
return awkward.concatenate(arrays, axis=axis)
def lepton_pairing(
Electron : List[ak_array],
Muon : List[ak_array]) -> (List[float], List[float], List[float]) :
"""
Electron and Muon contains values of ["PT", "Eta", "Phi", "Charge"] respectively
"""
return_Electron = ak.Array([])
return_Muon = ak.Array([])
delta_etas = []
delta_phis = []
Electron = ak.to_list(Electron)
Muon = ak.to_list(Muon)
n_events = len(Electron[0]) # each particle has same n of events but may have
# different n of particles in each event
for idx in range(n_events): # for each event
for electron_jdx in range(len(Electron[0][idx])):
for muon_jdx in range(len(Muon[0][idx])):
print("Electron[-1][idx][electron_jdx]: ", Electron[-1][idx][electron_jdx])
print("Muon[-1][idx][muon_jdx]: ", Muon[-1][idx][muon_jdx])
if (Electron[-1][idx][electron_jdx] == -Muon[-1][idx][muon_jdx] ):
# if charge sign of electron is opposite of muon, then pairing is made and
# add the relevant values to the final return value
e_placeholder = []
for value in Electron:
e_placeholder.append(value[idx][electron_jdx])
return_Electron = ak.concatenate((return_Electron, ak.Array([e_placeholder])), axis =0)
# print(ak.to_numpy(return_Electron).shape)
# print("return_Electron: ", return_Electron)
m_placeholder = []
for value in Muon:
m_placeholder.append(value[idx][muon_jdx])
return_Muon = ak.concatenate((return_Muon, ak.Array([m_placeholder])), axis =0)
# print(ak.to_numpy(return_Muon).shape)
# print("return_Muon: ", return_Muon)
# add delta eta and phi values
delta_etas.append(delta_eta(Electron[1][idx][electron_jdx], Muon[1][idx][muon_jdx]))
delta_phis.append(delta_phi(Electron[-2][idx][electron_jdx], Muon[-2][idx][muon_jdx]))
# assign the respective charges as 10 (arbitrary number) to signify that they
# have been chosen already, and don't get detected next time
# print(type(Electron[-1][idx][electron_jdx]))
Electron[-1][idx][electron_jdx] = 10.0
Muon[-1][idx][muon_jdx] = 10.0
delta_package = np.array([delta_etas, delta_phis])
return (ak.to_numpy(return_Electron), ak.to_numpy(return_Muon), delta_package)
def __call__(self, *args, **kwargs):
lens = set(map(len, args))
if len(lens) != 1:
raise ValueError("inconsistent *args len")
return awkward.concatenate([
self.func(*(a[off:off + self.chunksize] for a in args), **kwargs)
for off in range(0, max(lens), self.chunksize)
])
def test_merge_parameters():
one = ak.from_iter([[121, 117, 99, 107, 121], [115, 116, 117, 102, 102]],
highlevel=False)
two = ak.from_iter(["good", "stuff"], highlevel=False)
assert ak.to_list(ak.concatenate([one, two])) == [
[121, 117, 99, 107, 121],
[115, 116, 117, 102, 102],
"good",
"stuff",
]
assert ak.to_list(ak.concatenate([two, one])) == [
"good",
"stuff",
[121, 117, 99, 107, 121],
[115, 116, 117, 102, 102],
]
def _getMatchIdxs(self, tree):
#no split here
truthMom = self._readArray(tree,"MergedSimCluster_boundaryEnergy")
truthEta = self._readArray(tree,"MergedSimCluster_impactPoint_eta")
truthPhi = self._readArray(tree,"MergedSimCluster_impactPoint_phi")
truthpos = ak1.concatenate([self._expand(truthEta),self._expand(truthPhi)],axis=-1)
impactEta = self._readArray(tree,"Track_HGCFront_eta")
impactPhi = self._readArray(tree,"Track_HGCFront_phi")
impactpos = ak1.concatenate([self._expand(impactEta),self._expand(impactPhi)],axis=-1)
trackPt = self._readArray(tree,"Track_pt")
trackVertEta = self._readArray(tree,"Track_eta")
trackMom = trackPt * np.cosh(trackVertEta)
#match by x,y, and momentum
finalidxs = []
for tpos, ipos, tmom, imom, ipt in zip(truthpos, impactpos, truthMom, trackMom, trackPt):
# create default
tpos, ipos, tmom, imom,ipt = tpos.to_numpy(), ipos.to_numpy(), tmom.to_numpy(), imom.to_numpy(), ipt.to_numpy()
tpos = np.expand_dims(tpos, axis=0) #one is truth
tmom = np.expand_dims(tmom, axis=0) #one is truth
ipos = np.expand_dims(ipos, axis=1)
imom = np.expand_dims(imom, axis=1)
ipt = np.expand_dims(ipt,axis=1)
#this is in cm.
posdiffsq = np.sum( (tpos[:,:,0:1]-ipos[:,:,0:1])**2 +deltaPhi(tpos[:,:,1:2],ipos[:,:,1:2])**2, axis=-1) # Trk x K
#this is in %
momdiff = 100.*np.abs(tmom - imom)/(imom+1e-3) #rel diff
#scale position by 100 (DeltaR)
totaldiff = np.sqrt(100.**2*posdiffsq + (momdiff*np.exp(-0.05*ipt))**2)#weight momentum difference less with higher momenta
closestSC = np.argmin(totaldiff, axis=1) # Trk
#more than 5 percent/1cm total difference
closestSC[totaldiff[np.arange(len(closestSC)),closestSC] > 5] = -1
finalidxs.append(closestSC)
return ak1.from_iter(finalidxs)
def test_empty_arrays_cartesian():
one = ak.Array([])
two = one = ak.Array([])
with pytest.raises(ValueError) as err:
ak.to_list(ak.cartesian([one, two]))
assert isinstance(err.value, ValueError)
ak.to_list(ak.concatenate([one, two], axis=0))
def test_list_array_concatenate():
one = ak.Array([[1, 2, 3], [], [4, 5]]).layout
two = ak.Array([[1.1, 2.2], [3.3, 4.4], [5.5]]).layout
one = ak.layout.ListArray64(one.starts, one.stops, one.content)
two = ak.layout.ListArray64(two.starts, two.stops, two.content)
assert ak.to_list(ak.concatenate([one, two], 0)) == [
[1, 2, 3],
[],
[4, 5],
[1.1, 2.2],
[3.3, 4.4],
[5.5],
]
assert ak.to_list(ak.concatenate([one, two], 1)) == [
[1, 2, 3, 1.1, 2.2],
[3.3, 4.4],
[4, 5, 5.5],
]
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 __iadd__(self, other):
for branch, branch_data in other.data.items():
if branch in self.data.keys():
if ak.count(self.data[branch], axis=None) == 0:
self.data[branch] = branch_data
else:
self.data[branch] = ak.concatenate(
[self.data[branch], branch_data])
else:
self.data[branch] = branch_data
return self
def _embed_subjets(events):
sj1 = events.fatjets.subJetIdx1.content
sj2 = events.fatjets.subJetIdx2.content
sjcontent = ak.concatenate([
ak.JaggedArray.fromoffsets(np.arange(len(sj1) + 1), sj1),
ak.JaggedArray.fromoffsets(np.arange(len(sj2) + 1), sj2),
],
axis=1)
subjetidx = events.fatjets.copy(content=sjcontent)
subjetidx = subjetidx[subjetidx >= 0] + events.subjets.starts
events.fatjets['subjets'] = subjetidx.copy(content=subjetidx.content.copy(
content=events.subjets.content[subjetidx.flatten().flatten()]))
def loop_sample(self, sample, info):
if sample == "Data":
data = True
else:
data = False
sel_evts = []
process_id = info["process_id"]
for year, year_info in info.items():
if year not in self.years:
continue
files = []
for path in year_info["paths"]:
files += glob.glob(path + "/*.root")
if len(files) == 0:
if self.debug > 0:
print(
"[LoopHelper] Sample %s, year %s, has 0 input files, skipping."
% (sample, year))
continue
counter = 0
for file in files:
counter += 1
if self.fast and counter >= 2:
continue
if self.debug > 0:
print("[LoopHelper] Loading file %s" % file)
events = self.load_file(file, data=data)
selected_events = self.get_mask(
events) #self.select_events(events)
selected_events["process_id"] = numpy.ones(
len(selected_events)) * process_id
if data:
selected_events["weight"] = numpy.ones(
len(selected_events))
else:
selected_events[
"weight"] = selected_events.genWeight * year_info[
"metadata"]["scale1fb"]
selected_events_trimmed = self.trim_events(
selected_events, data)
sel_evts.append(selected_events_trimmed)
selected_events_full = awkward.concatenate(sel_evts)
return selected_events_full
def test_records_concatenate():
one = ak.Array([{
"x": 1,
"y": [1]
}, {
"x": 2,
"y": [1, 2]
}, {
"x": 3,
"y": [1, 2, 3]
}]).layout
two = ak.Array([{"y": [], "x": 4}, {"y": [3, 2, 1], "x": 5}]).layout
assert ak.to_list(ak.concatenate([one, two], 0)) == [
{
"x": 1,
"y": [1]
},
{
"x": 2,
"y": [1, 2]
},
{
"x": 3,
"y": [1, 2, 3]
},
{
"y": [],
"x": 4
},
{
"y": [3, 2, 1],
"x": 5
},
]
with pytest.raises(ValueError):
ak.to_list(ak.concatenate([one, two], 1))
with pytest.raises(ValueError):
ak.to_list(ak.concatenate([one, two], 2))
