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))