def run_deltar_matching(store, target, drname='deltaR', radius=0.4, unique=False, sort=False): """ Running a delta R matching of some object collection "store" of dimension NxS with some target collection "target" of dimension NxT, The return object will have dimension NxSxT' where objects in the T' contain all "target" objects within the delta R radius. The delta R between the store and target object will be stored in the field `deltaR`. If the unique flag is turned on, then objects in the target collection will only be associated to the closest object. If the sort flag is turned on, then the target collection will be sorted according to the computed `deltaR`. """ _, target = ak.unzip(ak.cartesian([store.eta, target], nested=True)) target[drname] = delta_r(store, target) if unique: # Additional filtering t_index = ak.argmin(target[drname], axis=-2) s_index = ak.local_index(store.eta, axis=-1) _, t_index = ak.unzip(ak.cartesian([s_index, t_index], nested=True)) target = target[s_index == t_index] # Cutting on the computed delta R target = target[target[drname] < radius] # Sorting according to the computed delta R if sort: idx = ak.argsort(target[drname], axis=-1) target = target[idx] return target
def get_taus(self, apply_selection=True): events = self.get_events() taus_dict = {"e": events.tau_e, "pt": events.tau_pt, "eta": events.tau_eta, "phi": events.tau_phi, "looseIsoAbs": events.tau_looseIsoAbs, "looseIsoRel": events.tau_looseIsoRel, "mediumIsoAbs": events.tau_mediumIsoAbs, "mediumIsoRel": events.tau_mediumIsoRel, "tightIsoAbs": events.tau_tightIsoAbs, "tightIsoRel": events.tau_tightIsoRel, "gen_e": events.gen_tau_e, "gen_pt": events.gen_tau_pt, "gen_eta": events.gen_tau_eta, "gen_phi": events.gen_tau_phi, "lepton_gen_match": events.lepton_gen_match, "deepTau_VSjet": events.deepTau_VSjet} if self.is_old: taus = ak.zip(taus_dict) index = ak.argsort(taus.pt, ascending=False) taus = taus[index] tau_1, tau_2 = ak.unzip(ak.combinations(taus, 2, axis=1)) else: taus_dict["vz"] = events.tau_vz taus = ak.zip(taus_dict) index = ak.argsort(taus.pt, ascending=False) taus = taus[index] tau_1, tau_2 = ak.unzip(ak.combinations(taus, 2, axis=1)) if apply_selection: L1taus = ak.zip({"e": events.L1tau_e, "pt": events.L1tau_pt, "eta": events.L1tau_eta, "phi": events.L1tau_phi}) # apply L1seed correction in case Pt28 and Pt30 seeds are considered L1taus, taus = L1seed_correction(L1taus, taus) # match taus with L1 taus taus = L1THLTTauMatching(L1taus, taus) tau_1, tau_2 = HLTJetPairDzMatchFilter(taus) # Return all possible pairs of tau which pass preselection return tau_1, tau_2
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 metric_table(self, other, axis=1, metric=lambda a, b: a.delta_r(b), return_combinations=False): """Return a list of a metric evaluated between this object and another. The two arrays should be broadcast-compatible on all axes other than the specified axis, which will be used to form a cartesian product. If axis=None, broadcast arrays directly. The return shape will be that of ``self`` with a new axis with shape of ``other`` appended at the specified axis depths. Parameters ---------- other : awkward.Array Another array with same shape in all but ``axis`` axis : int, optional The axis to form the cartesian product (default 1). If None, the metric is directly evaluated on the input arrays (i.e. they should broadcast) metric : callable A function of two arguments, returning a scalar. The default metric is `delta_r`. return_combinations : bool If True return the combinations of inputs as well as an unzipped tuple """ if axis is None: a, b = self, other else: a, b = awkward.unzip( awkward.cartesian([self, other], axis=axis, nested=True)) mval = metric(a, b) if return_combinations: return mval, (a, b) return mval
def get_gen_taus(self): events = self.get_gen_events() gen_taus = ak.zip({"gen_e": events.gen_tau_e, "gen_pt": events.gen_tau_pt, "gen_eta": events.gen_tau_eta, "gen_phi": events.gen_tau_phi, "lepton_gen_match": events.lepton_gen_match}) index = ak.argsort(gen_taus.gen_pt, ascending=False) gen_taus = gen_taus[index] gen_tau_1, gen_tau_2 = ak.unzip(ak.combinations(gen_taus, 2, axis=1)) return gen_tau_1, gen_tau_2
def L1THLTTauMatching(L1taus, taus): dR_matching = 0.5 tau_inpair, L1_inpair = ak.unzip(ak.cartesian([taus, L1taus], nested=True)) # dR = delta_r(L1_inpair, tau_inpair) dR = delta_r(tau_inpair, L1_inpair) # # print(dR[range(2)]) # # consider only L1taus for which there is at least 1 matched tau # tau_inpair = tau_inpair[dR<dR_matching] mask = ak.sum(dR < dR_matching, axis=-1) > 0 # tau_inpair = tau_inpair[mask] # # take first matched tau for each L1tau # L2taus = tau_inpair[:,:,0] # # print(L2taus[range(2)]) L2taus = taus[mask] return L2taus
def HLTJetPairDzMatchFilter(L2taus): jetMinPt = 20.0 jetMaxEta = 2.1 jetMinDR = 0.5 jetMaxDZ = 0.2 L2taus = L2taus[reco_tau_selection(L2taus, minPt=jetMinPt, maxEta=jetMaxEta)] # Take all possible pairs of L2 taus L2tau_1, L2tau_2 = ak.unzip(ak.combinations(L2taus, 2, axis=1)) dr2 = delta_r2(L2tau_1, L2tau_2) dz = delta_z(L2tau_1, L2tau_2) pair_mask = (dr2 >= jetMinDR * jetMinDR) & (abs(dz) <= jetMaxDZ) # ev_mask = ak.sum(pair_mask, axis=1) > 0 return L2tau_1[pair_mask], L2tau_2[pair_mask]
def eval_model(model, dataset, config, outdir): ibatch = 0 jetdef = fastjet.JetDefinition(fastjet.antikt_algorithm, 0.4) for elem in tqdm(dataset, desc="Evaluating model"): y_pred = model.predict(elem["X"], verbose=False) np_outfile = "{}/pred_batch{}.npz".format(outdir, ibatch) ygen = unpack_target(elem["ygen"], config["dataset"]["num_output_classes"], config) ycand = unpack_target(elem["ycand"], config["dataset"]["num_output_classes"], config) outs = {} for key in y_pred.keys(): outs["gen_{}".format(key)] = ygen[key].numpy() outs["cand_{}".format(key)] = ycand[key].numpy() outs["pred_{}".format(key)] = y_pred[key] jets_coll = {} jets_const = {} for typ in ["gen", "cand", "pred"]: cls_id = np.argmax(outs["{}_cls".format(typ)], axis=-1) valid = cls_id != 0 pt = awkward.from_iter( [y[m][:, 0] for y, m in zip(outs["{}_pt".format(typ)], valid)]) eta = awkward.from_iter([ y[m][:, 0] for y, m in zip(outs["{}_eta".format(typ)], valid) ]) phi = np.arctan2(outs["{}_sin_phi".format(typ)], outs["{}_cos_phi".format(typ)]) phi = awkward.from_iter([y[m][:, 0] for y, m in zip(phi, valid)]) e = awkward.from_iter([ y[m][:, 0] for y, m in zip(outs["{}_energy".format(typ)], valid) ]) vec = vector.arr({"pt": pt, "eta": eta, "phi": phi, "e": e}) cluster = fastjet.ClusterSequence(vec.to_xyzt(), jetdef) jets = cluster.inclusive_jets() jet_constituents = cluster.constituent_index() jets_coll[typ] = jets[jets.pt > 5.0] jets_const[typ] = jet_constituents[jets.pt > 5.0] for key in ["pt", "eta", "phi", "energy"]: outs["jets_gen_{}".format(key)] = awkward.to_numpy( awkward.flatten(getattr(jets_coll["gen"], key))) outs["jets_cand_{}".format(key)] = awkward.to_numpy( awkward.flatten(getattr(jets_coll["cand"], key))) outs["jets_pred_{}".format(key)] = awkward.to_numpy( awkward.flatten(getattr(jets_coll["pred"], key))) # DeltaR match between genjets and PF/MLPF jets cart = awkward.cartesian([jets_coll["gen"], jets_coll["pred"]], nested=True) jets_a, jets_b = awkward.unzip(cart) drs = deltar(jets_a, jets_b) match_gen_to_pred = [awkward.where(d < 0.1) for d in drs] m0 = awkward.from_iter([m[0] for m in match_gen_to_pred]) m1 = awkward.from_iter([m[1] for m in match_gen_to_pred]) j1s = jets_coll["gen"][m0] j2s = jets_coll["pred"][m1] outs["jets_pt_gen_to_pred"] = np.stack([ awkward.to_numpy(awkward.flatten(j1s.pt)), awkward.to_numpy(awkward.flatten(j2s.pt)) ], axis=-1) cart = awkward.cartesian([jets_coll["gen"], jets_coll["cand"]], nested=True) jets_a, jets_b = awkward.unzip(cart) drs = deltar(jets_a, jets_b) match_gen_to_pred = [awkward.where(d < 0.1) for d in drs] m0 = awkward.from_iter([m[0] for m in match_gen_to_pred]) m1 = awkward.from_iter([m[1] for m in match_gen_to_pred]) j1s = jets_coll["gen"][m0] j2s = jets_coll["cand"][m1] outs["jets_pt_gen_to_cand"] = np.stack([ awkward.to_numpy(awkward.flatten(j1s.pt)), awkward.to_numpy(awkward.flatten(j2s.pt)) ], axis=-1) np.savez(np_outfile, X=elem["X"], **outs) ibatch += 1
def parse_to_parquet(base_output_filename: Union[Path, str], store_only_necessary_columns: bool, input_filename: Union[Path, str], events_per_chunk: int, parser: str = "pandas", max_chunks: int = -1, compression: str = "zstd", compression_level: Optional[int] = None) -> None: """ Parse the JETSCAPE ASCII and convert it to parquet, (potentially) storing only the minimum necessary columns. Args: base_output_filename: Basic output filename. Should include the entire path. store_only_necessary_columns: If True, store only the necessary columns, rather than all of them. input_filename: Filename of the input JETSCAPE ASCII file. events_per_chunk: Number of events to be read per chunk. parser: Name of the parser. Default: "pandas". max_chunks: Maximum number of chunks to read. Default: -1. compression: Compression algorithm for parquet. Default: "zstd". Options include: ["snappy", "gzip", "ztsd"]. "gzip" is slightly better for storage, but slower. See the compression tests and parquet docs for more. compression_level: Compression level for parquet. Default: `None`, which lets parquet choose the best value. Returns: None. The parsed events are stored in parquet files. """ # Validation base_output_filename = Path(base_output_filename) # Setup the base output directory base_output_filename.parent.mkdir(parents=True, exist_ok=True) # We will check which fields actually exist when writing. possible_fields_containing_floats = ["event_plane_angle", "event_weight", "cross_section", "cross_section_error", "px", "py", "pz", "E"] for i, arrays in enumerate(read(filename=input_filename, events_per_chunk=events_per_chunk, parser=parser)): # Reduce to the minimum required data. if store_only_necessary_columns: arrays = full_events_to_only_necessary_columns_E_px_py_pz(arrays) else: # To match the steps taken when reducing the columns, we'll re-zip with the depth limited to 1. # As of April 2021, I'm not certainly this is truly required anymore, but it may be needed for # parquet writing to be successful (apparently parquet couldn't handle lists of structs sometime # in 2020. The status in April 2021 is unclear, but not worth digging into now). arrays = ak.zip( dict(zip(ak.fields(arrays), ak.unzip(arrays))), depth_limit = 1 ) # If converting in chunks, add an index to the output file so the chunks don't overwrite each other. if events_per_chunk > 0: suffix = base_output_filename.suffix output_filename = (base_output_filename.parent / f"{base_output_filename.stem}_{i:02}").with_suffix(suffix) else: output_filename = base_output_filename # Optimize the output # Additional parquet options are based on https://stackoverflow.com/a/66854439/12907985 # byte_stream_fields apparently only work for float fields. Other fields should be handled # by use_dictionary. Apparently it can't handle this automatically, we so we have to define it # ourselves. This is a bit brittle if fields change, but they don't change so often, and # it's simpler than parsing field types, so it should be fine for now. byte_stream_fields = [field for field in ak.fields(arrays) if field in possible_fields_containing_floats] dict_fields = [field for field in ak.fields(arrays) if field not in possible_fields_containing_floats] # logger.debug(f"dict_fields: {dict_fields}") # logger.debug(f"byte_stream_fields: {byte_stream_fields}") # Parquet with zlib seems to do about the same as ascii tar.gz when we drop unneeded columns. # And it should load much faster! ak.to_parquet( arrays, output_filename, compression=compression, compression_level=compression_level, explode_records=False, # Additional parquet options are based on https://stackoverflow.com/a/66854439/12907985 #use_dictionary=True, #use_byte_stream_split=True, use_dictionary=dict_fields, use_byte_stream_split=byte_stream_fields, ) # Break now so we don't have to read the next chunk. if (i + 1) == max_chunks: break
def process(self, events): output = self._accumulator.identity() jets=events.Jet jetSel = (jets.pt>30) & (abs(jets.eta)<2.4) tightJet = jets[jetSel] bJet = tightJet[tightJet.btagDeepFlavB > 0.642] muons = events.Muon muonSel = (muons.pt>30) & (abs(muons.eta)<2.4) tightMuon = muons[muonSel] ele = events.Electron eleSel = (ele.pt>35)&(abs(ele.eta)<2.4) tightEle = ele[eleSel] eventSel = (((ak.num(tightMuon)==1) | (ak.num(tightEle)==1)) & (ak.num(tightJet)>= 3) & (ak.num(bJet)>=1) ) final = events[eventSel] #####GENPART MATCHING ###### genPart = final.GenPart tops = genPart[abs(genPart.pdgId)==6] #The isLastCopy Flag filters out copy Genparticles: tops = tops[tops.hasFlags('isLastCopy')] tDecay = tops.distinctChildren tDecay = tDecay[tDecay.hasFlags('isLastCopy')] t_Events=tDecay[abs(tDecay.pdgId)==5] W = tDecay[abs(tDecay.pdgId)==24] W = W[W.hasFlags('isLastCopy')] WDecay = W.distinctChildren WDecay = WDecay[WDecay.hasFlags('isLastCopy')] #t_events is the lone bottom, W_events is the -> two jets #select the hadronically decaying W W_Events=ak.flatten(WDecay[ak.all(abs(WDecay.pdgId)<=8,axis=-1)],axis=3) #print(qqb) #HadW is mask for Quark deacying W boson hadW = ak.num(W_Events,axis=2)==2 #filters out t_events that have a hadronically decayign W Boson hadB = t_Events[hadW] hadB = ak.flatten(hadB,axis=2) W_quarks = W_Events[hadW] W_quarks = ak.flatten(W_quarks,axis=2) #concatentating these two arrays make an array of events with the correctly decaying GenParticles. qqb = ak.concatenate([hadB,W_quarks],axis=1) #####GEN JET MATCHING ###### 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] #Match Gen part to gen jet matchedGenJets=qqb.nearest(final.GenJet) #match gen to reco matchedJets=matchedGenJets.nearest(finalJets) ### VALIDATION ### 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) jetcombos=ak.flatten(combs) j1,j2,j3=ak.unzip(jetcombos) output["dR12"]+=processor.column_accumulator(ak.to_numpy(j1.delta_r(j2))) output["dR13"]+=processor.column_accumulator(ak.to_numpy(j1.delta_r(j3))) output["dR23"]+=processor.column_accumulator(ak.to_numpy(j2.delta_r(j3))) output["j1btag"]+=processor.column_accumulator(ak.to_numpy(j1.btagCSVV2)) output["j2btag"]+=processor.column_accumulator(ak.to_numpy(j1.btagCSVV2)) output["j3btag"]+=processor.column_accumulator(ak.to_numpy(j1.btagCSVV2)) output["j1area"]+=processor.column_accumulator(ak.to_numpy(j1.area)) output["j2area"]+=processor.column_accumulator(ak.to_numpy(j2.area)) output["j3area"]+=processor.column_accumulator(ak.to_numpy(j3.area)) output["j12deta"]+=processor.column_accumulator(ak.to_numpy(j1.eta-j2.eta)) output["j23deta"]+=processor.column_accumulator(ak.to_numpy(j2.eta-j3.eta)) output["j13deta"]+=processor.column_accumulator(ak.to_numpy(j1.eta-j3.eta)) output["j12dphi"]+=processor.column_accumulator(ak.to_numpy(j1.phi-j2.phi)) output["j23dphi"]+=processor.column_accumulator(ak.to_numpy(j2.phi-j3.phi)) output["j13dphi"]+=processor.column_accumulator(ak.to_numpy(j1.phi-j3.phi)) output["j1j2mass"]+=processor.column_accumulator(ak.to_numpy(j1.mass+j2.mass)) output["j2j3mass"]+=processor.column_accumulator(ak.to_numpy(j2.mass+j3.mass)) output["j1j3mass"]+=processor.column_accumulator(ak.to_numpy(j1.mass+j3.mass)) output["j1pt"]+=processor.column_accumulator(ak.to_numpy(j1.pt)) output["j1phi"]+=processor.column_accumulator(ak.to_numpy(j1.phi)) output["j1eta"]+=processor.column_accumulator(ak.to_numpy(abs(j1.eta))) output["j1mass"]+=processor.column_accumulator(ak.to_numpy(j1.mass)) output["j2pt"]+=processor.column_accumulator(ak.to_numpy(j2.pt)) output["j2phi"]+=processor.column_accumulator(ak.to_numpy(j2.phi)) output["j2eta"]+=processor.column_accumulator(ak.to_numpy(abs(j2.eta))) output["j2mass"]+=processor.column_accumulator(ak.to_numpy(j2.mass)) output["j3pt"]+=processor.column_accumulator(ak.to_numpy(j3.pt)) output["j3phi"]+=processor.column_accumulator(ak.to_numpy(j3.phi)) output["j3eta"]+=processor.column_accumulator(ak.to_numpy(abs(j3.eta))) output["j3mass"]+=processor.column_accumulator(ak.to_numpy(j3.mass)) output["event"]+=processor.column_accumulator(ak.to_numpy(ak.flatten(ak.broadcast_arrays(final.event,combs['0'].pt)[0]))) output["truth"]+=processor.column_accumulator(ak.to_numpy(trutharray).astype(int)) return output
Kst_parent_cut_pi = abs(Kst["pi"]["pdg"]) == abs(lp.Kst_892_0.pdgid) Kst_parent_cut_k = abs(Kst["k"]["pdg"]) == abs(lp.Kst_892_0.pdgid) Kst_parent_cut_all = np.logical_and(Kst_parent_cut_pi, Kst_parent_cut_k) Kst = Kst[Kst_parent_cut_all] #Ensure kaon and pion come from the same K* Kst_parent_cut_kpi = Kst["pi","MC_parentindex"] == Kst["k","MC_parentindex"] Kst = Kst[Kst_parent_cut_kpi] for comp in ["x","y","z"]: Kst[f"p{comp}"] = Kst["k",f"p{comp}"] + Kst["pi",f"p{comp}"] Kst["p"] = np.sqrt(Kst["px"]**2 + Kst["py"]**2 + Kst["pz"]**2) Kst["e"] = np.sqrt(Kst["p"]**2 + Kst["mass"]**2) tau = ak.combinations(pions,3) tau["pi_1"], tau["pi_2"], tau["pi_3"] = ak.unzip(tau) #Charge cut on the pions tau_charge_cut_1 = np.sign(tau["pi_1","charge"]) != np.sign(tau["pi_3","charge"]) tau_charge_cut_2 = np.sign(tau["pi_2","charge"]) != np.sign(tau["pi_3","charge"]) tau_charge_cut = np.logical_and(tau_charge_cut_1, tau_charge_cut_2) tau = tau[tau_charge_cut] PDG_pi_m = lp.pi_plus.mass/1000. tau["mass"] = kinematics_flat.mass([tau["pi_1"], tau["pi_2"], tau["pi_3"]], [PDG_pi_m, PDG_pi_m, PDG_pi_m]) tau_m_low_cut_val = 3*PDG_pi_m tau_m_high_cut_val = lp.tau_plus.mass/1000. tau_m_low_cut = tau["mass"] >= tau_m_low_cut_val tau_m_high_cut = tau["mass"] <= lp.tau_plus.mass/1000. tau_m_cut = np.logical_and(tau_m_low_cut, tau_m_high_cut)
def test_zip(): x = ak.Array([[1, 2, 3], [], [4, 5], [6], [7, 8, 9, 10]]) y = ak.Array([1.1, 2.2, 3.3, 4.4, 5.5]) one = ak.zip({"x": x, "y": y}) two = ak.zip({"x": x, "y": y}, depth_limit=1) xx, yy = ak.unzip(two) assert isinstance(one.layout, ak.layout.Content) assert isinstance(two.layout, ak.layout.Content) assert isinstance(xx.layout, ak.layout.Content) assert isinstance(yy.layout, ak.layout.Content) assert ak.to_list(one) == [ [{ "x": 1, "y": 1.1 }, { "x": 2, "y": 1.1 }, { "x": 3, "y": 1.1 }], [], [{ "x": 4, "y": 3.3 }, { "x": 5, "y": 3.3 }], [{ "x": 6, "y": 4.4 }], [ { "x": 7, "y": 5.5 }, { "x": 8, "y": 5.5 }, { "x": 9, "y": 5.5 }, { "x": 10, "y": 5.5 }, ], ] assert ak.to_list(two) == [ { "x": [1, 2, 3], "y": 1.1 }, { "x": [], "y": 2.2 }, { "x": [4, 5], "y": 3.3 }, { "x": [6], "y": 4.4 }, { "x": [7, 8, 9, 10], "y": 5.5 }, ] if not ak._util.py27 and not ak._util.py35: assert ak.to_list(xx) == [[1, 2, 3], [], [4, 5], [6], [7, 8, 9, 10]] assert ak.to_list(yy) == [1.1, 2.2, 3.3, 4.4, 5.5] x = ak.repartition(x, 3) assert isinstance(x.layout, ak.partition.PartitionedArray) assert ak.to_list(x) == [[1, 2, 3], [], [4, 5], [6], [7, 8, 9, 10]] one = ak.zip({"x": x, "y": y}) two = ak.zip({"x": x, "y": y}, depth_limit=1) xx, yy = ak.unzip(two) assert isinstance(one.layout, ak.partition.PartitionedArray) assert isinstance(two.layout, ak.partition.PartitionedArray) assert isinstance(xx.layout, ak.partition.PartitionedArray) assert isinstance(yy.layout, ak.partition.PartitionedArray) assert ak.to_list(one) == [ [{ "x": 1, "y": 1.1 }, { "x": 2, "y": 1.1 }, { "x": 3, "y": 1.1 }], [], [{ "x": 4, "y": 3.3 }, { "x": 5, "y": 3.3 }], [{ "x": 6, "y": 4.4 }], [ { "x": 7, "y": 5.5 }, { "x": 8, "y": 5.5 }, { "x": 9, "y": 5.5 }, { "x": 10, "y": 5.5 }, ], ] assert ak.to_list(two) == [ { "x": [1, 2, 3], "y": 1.1 }, { "x": [], "y": 2.2 }, { "x": [4, 5], "y": 3.3 }, { "x": [6], "y": 4.4 }, { "x": [7, 8, 9, 10], "y": 5.5 }, ] if not ak._util.py27 and not ak._util.py35: assert ak.to_list(xx) == [[1, 2, 3], [], [4, 5], [6], [7, 8, 9, 10]] assert ak.to_list(yy) == [1.1, 2.2, 3.3, 4.4, 5.5] y = ak.repartition(y, 2) assert isinstance(x.layout, ak.partition.PartitionedArray) assert ak.to_list(y) == [1.1, 2.2, 3.3, 4.4, 5.5] one = ak.zip({"x": x, "y": y}) two = ak.zip({"x": x, "y": y}, depth_limit=1) xx, yy = ak.unzip(two) assert isinstance(one.layout, ak.partition.PartitionedArray) assert isinstance(two.layout, ak.partition.PartitionedArray) assert isinstance(xx.layout, ak.partition.PartitionedArray) assert isinstance(yy.layout, ak.partition.PartitionedArray) assert ak.to_list(one) == [ [{ "x": 1, "y": 1.1 }, { "x": 2, "y": 1.1 }, { "x": 3, "y": 1.1 }], [], [{ "x": 4, "y": 3.3 }, { "x": 5, "y": 3.3 }], [{ "x": 6, "y": 4.4 }], [ { "x": 7, "y": 5.5 }, { "x": 8, "y": 5.5 }, { "x": 9, "y": 5.5 }, { "x": 10, "y": 5.5 }, ], ] assert ak.to_list(two) == [ { "x": [1, 2, 3], "y": 1.1 }, { "x": [], "y": 2.2 }, { "x": [4, 5], "y": 3.3 }, { "x": [6], "y": 4.4 }, { "x": [7, 8, 9, 10], "y": 5.5 }, ] if not ak._util.py27 and not ak._util.py35: assert ak.to_list(xx) == [[1, 2, 3], [], [4, 5], [6], [7, 8, 9, 10]] assert ak.to_list(yy) == [1.1, 2.2, 3.3, 4.4, 5.5] x = ak.repartition(x, None) assert isinstance(x.layout, ak.layout.Content) assert ak.to_list(x) == [[1, 2, 3], [], [4, 5], [6], [7, 8, 9, 10]] one = ak.zip({"x": x, "y": y}) two = ak.zip({"x": x, "y": y}, depth_limit=1) xx, yy = ak.unzip(two) assert isinstance(one.layout, ak.partition.PartitionedArray) assert isinstance(two.layout, ak.partition.PartitionedArray) assert isinstance(xx.layout, ak.partition.PartitionedArray) assert isinstance(yy.layout, ak.partition.PartitionedArray) assert ak.to_list(one) == [ [{ "x": 1, "y": 1.1 }, { "x": 2, "y": 1.1 }, { "x": 3, "y": 1.1 }], [], [{ "x": 4, "y": 3.3 }, { "x": 5, "y": 3.3 }], [{ "x": 6, "y": 4.4 }], [ { "x": 7, "y": 5.5 }, { "x": 8, "y": 5.5 }, { "x": 9, "y": 5.5 }, { "x": 10, "y": 5.5 }, ], ] assert ak.to_list(two) == [ { "x": [1, 2, 3], "y": 1.1 }, { "x": [], "y": 2.2 }, { "x": [4, 5], "y": 3.3 }, { "x": [6], "y": 4.4 }, { "x": [7, 8, 9, 10], "y": 5.5 }, ] if not ak._util.py27 and not ak._util.py35: assert ak.to_list(xx) == [[1, 2, 3], [], [4, 5], [6], [7, 8, 9, 10]] assert ak.to_list(yy) == [1.1, 2.2, 3.3, 4.4, 5.5] y = ak.repartition(y, None) assert isinstance(y.layout, ak.layout.Content) assert ak.to_list(y) == [1.1, 2.2, 3.3, 4.4, 5.5] one = ak.zip({"x": x, "y": y}) two = ak.zip({"x": x, "y": y}, depth_limit=1) xx, yy = ak.unzip(two) assert isinstance(one.layout, ak.layout.Content) assert isinstance(two.layout, ak.layout.Content) assert isinstance(xx.layout, ak.layout.Content) assert isinstance(yy.layout, ak.layout.Content) assert ak.to_list(one) == [ [{ "x": 1, "y": 1.1 }, { "x": 2, "y": 1.1 }, { "x": 3, "y": 1.1 }], [], [{ "x": 4, "y": 3.3 }, { "x": 5, "y": 3.3 }], [{ "x": 6, "y": 4.4 }], [ { "x": 7, "y": 5.5 }, { "x": 8, "y": 5.5 }, { "x": 9, "y": 5.5 }, { "x": 10, "y": 5.5 }, ], ] assert ak.to_list(two) == [ { "x": [1, 2, 3], "y": 1.1 }, { "x": [], "y": 2.2 }, { "x": [4, 5], "y": 3.3 }, { "x": [6], "y": 4.4 }, { "x": [7, 8, 9, 10], "y": 5.5 }, ] if not ak._util.py27 and not ak._util.py35: assert ak.to_list(xx) == [[1, 2, 3], [], [4, 5], [6], [7, 8, 9, 10]] assert ak.to_list(yy) == [1.1, 2.2, 3.3, 4.4, 5.5]
def run(B_type): #Load the MC file_name = f"p8_ee_Zbb_ecm91_EvtGen_{B_type}2TauNuTAUHADNU" file = uproot.open(f"{loc.IN}/{mode}/{file_name}.root") tree = file['events'] #Awkward array of the MC - filter event variables if running on Bu sample if (B_type == "Bc"): events = tree.arrays(library="ak", how="zip") #Keep a manageable subset of the Bu events (1M too many to handle in memory) else: events = tree.arrays(library="ak", how="zip", entry_stop=20000) #Reco particles rp = events["RP"] #Index cut to remove negative indices (particles produced before hadronisation) parentindex_cut = rp["MC_parentindex"] >= 0 rp = rp[parentindex_cut] grandparentindex_cut = rp["MC_grandparentindex"] >= 0 rp = rp[grandparentindex_cut] #Match to MC truth partners rp["parent_pdg"] = events["MC", "pdg"][rp["MC_parentindex"]] rp["grandparent_pdg"] = events["MC", "pdg"][rp["MC_grandparentindex"]] rp["pdg"] = events["MC", "pdg"][rp["MC_index"]] #Get the production vertex of the track (truth-level) for v in ["x", "y", "z"]: rp[f"vertex_{v}"] = events["MC", f"vertex_{v}"][rp["MC_index"]] #Get the production vertex of the parent of the tracks (for the pions, this will give the tau production vertex i.e. the Bc decay vertex) for v in ["x", "y", "z"]: rp[f"parent_vertex_{v}"] = events["MC", f"vertex_{v}"][rp["MC_parentindex"]] #Separation of vertex from PV (combined tau and Bc flight) rp["PVsep"] = np.sqrt(rp[f"vertex_x"]**2 + rp[f"vertex_y"]**2 + rp[f"vertex_z"]**2) #Separation of the Bc vertex from the PV rp["PVsep_parent"] = np.sqrt(rp[f"parent_vertex_x"]**2 + rp[f"parent_vertex_y"]**2 + rp[f"parent_vertex_z"]**2) #Charged tracks to make tau from, will truth match them to pi+/- ID below p_cut = rp["p"] > 0. pions = rp[p_cut] #Keep charged tracks charge_cut = abs(rp["charge"]) == 1 pions = pions[charge_cut] #Build the tau -> 3pi tau = ak.combinations(pions, 3) tau["pi_1"], tau["pi_2"], tau["pi_3"] = ak.unzip(tau) PDG_pi_m = lp.pi_plus.mass / 1000. tau["mass"] = kinematics_flat.mass([tau["pi_1"], tau["pi_2"], tau["pi_3"]], [PDG_pi_m, PDG_pi_m, PDG_pi_m]) tau_m_low_cut_val = 3 * PDG_pi_m tau_m_high_cut_val = lp.tau_plus.mass / 1000. tau_m_low_cut = tau["mass"] >= tau_m_low_cut_val tau_m_high_cut = tau["mass"] <= lp.tau_plus.mass / 1000. tau_m_cut = np.logical_and(tau_m_low_cut, tau_m_high_cut) tau = tau[tau_m_cut] #Truth-matching for signal #Pions are all pions pi1_cut = abs(tau["pi_1"]["pdg"]) == abs(lp.pi_plus.pdgid) pi2_cut = abs(tau["pi_2"]["pdg"]) == abs(lp.pi_plus.pdgid) pi3_cut = abs(tau["pi_3"]["pdg"]) == abs(lp.pi_plus.pdgid) pi12_cut = np.logical_and(pi1_cut, pi2_cut) pi_cut_all = np.logical_and(pi12_cut, pi3_cut) tau = tau[pi_cut_all] #Pion parents are all tau tau_parent_cut_pi1 = abs(tau["pi_1"]["parent_pdg"]) == abs( lp.tau_plus.pdgid) tau_parent_cut_pi2 = abs(tau["pi_2"]["parent_pdg"]) == abs( lp.tau_plus.pdgid) tau_parent_cut_pi3 = abs(tau["pi_3"]["parent_pdg"]) == abs( lp.tau_plus.pdgid) tau_parent_cut_12 = np.logical_and(tau_parent_cut_pi1, tau_parent_cut_pi2) tau_parent_cut_all = np.logical_and(tau_parent_cut_12, tau_parent_cut_pi3) tau = tau[tau_parent_cut_all] #Pion grandparent are all Bu / Bc parent_id = {"Bu": lp.B_plus.pdgid, "Bc": lp.B_c_plus.pdgid} tau_grandparent_cut_pi1 = abs(tau["pi_1"]["grandparent_pdg"]) == abs( parent_id[B_type]) tau_grandparent_cut_pi2 = abs(tau["pi_2"]["grandparent_pdg"]) == abs( parent_id[B_type]) tau_grandparent_cut_pi3 = abs(tau["pi_3"]["grandparent_pdg"]) == abs( parent_id[B_type]) tau_grandparent_cut_12 = np.logical_and(tau_grandparent_cut_pi1, tau_grandparent_cut_pi2) tau_grandparent_cut_all = np.logical_and(tau_grandparent_cut_12, tau_grandparent_cut_pi3) tau = tau[tau_grandparent_cut_all] #Ensure pions come from the same tau tau_parent_cut_12 = tau["pi_1", "MC_parentindex"] == tau["pi_2", "MC_parentindex"] tau_parent_cut_13 = tau["pi_1", "MC_parentindex"] == tau["pi_3", "MC_parentindex"] tau_parent_cut = np.logical_and(tau_parent_cut_12, tau_parent_cut_13) tau = tau[tau_parent_cut] #Net charge of the tau tau["charge"] = tau["pi_1", "charge"] + tau["pi_2", "charge"] + tau["pi_3", "charge"] #Tau energy (visible energy of the signal) for comp in ["x", "y", "z"]: tau[f"p{comp}"] = tau["pi_1", f"p{comp}"] + tau[ "pi_2", f"p{comp}"] + tau["pi_1", f"p{comp}"] tau["p"] = np.sqrt(tau["px"]**2 + tau["py"]**2 + tau["pz"]**2) tau["e"] = np.sqrt(tau["mass"]**2 + tau["p"]**2) #Tau flight and Bc flight from PV (using true vertices) tau["PVsep"] = tau["pi_1", "PVsep"] #x, y, and z coordinates of the tau vertex for v in ["x", "y", "z"]: tau[f"PVsep_{v}"] = tau["pi_1", f"vertex_{v}"] tau["PVsep_parent"] = tau["pi_1", "PVsep_parent"] #Thrust axis co-ordinates to persist for v in ["x", "y", "z"]: tau[f"EVT_thrust_{v}"] = events[f"EVT_thrust_{v}"] #Hemisphere energies for hem in ["0", "1"]: events[f"EVT_thrutshemis{hem}_e"] = events[ f"EVT_thrutshemis{hem}_echarged"] + events[ f"EVT_thrutshemis{hem}_eneutral"] #Total energy (sum of both hemispheres) events[ "EVT_e"] = events["EVT_thrutshemis0_e"] + events["EVT_thrutshemis1_e"] #Min and max hemisphere energies per-event events["EVT_thrutshemis_e_min"] = np.minimum(events["EVT_thrutshemis0_e"], events["EVT_thrutshemis1_e"]) events["EVT_thrutshemis_e_max"] = np.maximum(events["EVT_thrutshemis0_e"], events["EVT_thrutshemis1_e"]) #Difference in energy between hemispheres (max - min so always positive) events["EVT_thrutshemis_e_diff"] = events[ "EVT_thrutshemis_e_max"] - events["EVT_thrutshemis_e_min"] fig, ax = plt.subplots(figsize=(9, 9)) plt.hist(events["EVT_thrutshemis_e_min"], bins=50, range=(0, 60), color="crimson", histtype='step', linewidth=2, label="Lower energy hemisphere per-event") plt.hist(events["EVT_thrutshemis_e_max"], bins=50, range=(0, 60), color="dodgerblue", histtype='step', linewidth=2, label="Higher energy hemisphere per-event") ax.tick_params(axis='both', which='major', labelsize=25) plt.xlim(0, 60) plt.axvline(lp.Z_0.mass / 2000., color='k', linestyle='--', label="$m(Z^0)/2$") plt.xlabel("Hemisphere energy [GeV]", fontsize=30) plt.legend(fontsize=18, loc="upper left") ut.create_dir(loc.PLOTS) fig.savefig(f"{loc.PLOTS}/{B_type}2TauNu_min_max_hemisphere_E.pdf") fig, ax = plt.subplots(figsize=(9, 9)) plt.hist(events["EVT_thrutshemis_e_diff"], bins=30, range=(0, 50), color="k", histtype='step', linewidth=2) ax.tick_params(axis='both', which='major', labelsize=25) plt.xlim(0, 50) plt.xlabel("Hemisphere energy difference [GeV]", fontsize=30) fig.savefig(f"{loc.PLOTS}/{B_type}2TauNu_diff_hemisphere_E.pdf") #Store hemisphere energy info in tau container for writing out later tau["EVT_e"] = events["EVT_e"] tau["EVT_thrutshemis_e_min"] = events["EVT_thrutshemis_e_min"] tau["EVT_thrutshemis_e_max"] = events["EVT_thrutshemis_e_max"] tau["EVT_thrutshemis_e_diff"] = events["EVT_thrutshemis_e_diff"] tau["EVT_thrutshemis0_e"] = events["EVT_thrutshemis0_e"] tau["EVT_thrutshemis1_e"] = events["EVT_thrutshemis1_e"] #Nominal B energy = m(Z) - sum of all energy apart from signal tau["nominal_B_e"] = lp.Z_0.mass / 1000. - (tau["EVT_e"] - tau["e"]) fig, ax = plt.subplots(figsize=(9, 9)) plt.hist(ak.flatten(tau["nominal_B_e"]), range=(0, 80), bins=30, color="k", histtype='step', linewidth=2) ax.tick_params(axis='both', which='major', labelsize=25) plt.title( "$m(Z) - \\Sigma$(all visible E apart from true $\\tau \\to 3\\pi$)", fontsize=25) plt.xlim(0, 80) plt.xlabel("Nominal $B$ energy [GeV]", fontsize=30) fig.savefig(f"{loc.PLOTS}/{B_type}2TauNu_nominal_B_E.pdf") #Plot the tau mass for truth-matched signal fig, ax = plt.subplots(figsize=(9, 9)) plt.hist(ak.flatten(tau["mass"]), range=(tau_m_low_cut_val, tau_m_high_cut_val), bins=30, color="k", histtype='step', linewidth=2) ax.tick_params(axis='both', which='major', labelsize=25) plt.xlim(tau_m_low_cut_val, tau_m_high_cut_val) plt.xlabel("$m(3\\pi)$ [GeV]", fontsize=30) fig.savefig(f"{loc.PLOTS}/{B_type}2TauNu_3pi_M.pdf") #Look at charged and neutral multiplicities in the two hemispheres #Events where hemisphere 0 is the minimum energy events_hem0_min_cut = events["EVT_thrutshemis0_e"] == events[ "EVT_thrutshemis_e_min"] events_hem0_min = events[events_hem0_min_cut] #Events where hemisphere 1 is the minimum energy events_hem1_min_cut = events["EVT_thrutshemis1_e"] == events[ "EVT_thrutshemis_e_min"] events_hem1_min = events[events_hem1_min_cut] #Get the charged and neutral energy and multiplicity for ptype in ["charged", "neutral"]: for var in ["n", "e"]: events_hem0_min[f"{var}{ptype}_min"] = events_hem0_min[ f"EVT_thrutshemis0_{var}{ptype}"] events_hem0_min[f"{var}{ptype}_max"] = events_hem0_min[ f"EVT_thrutshemis1_{var}{ptype}"] events_hem1_min[f"{var}{ptype}_min"] = events_hem1_min[ f"EVT_thrutshemis1_{var}{ptype}"] events_hem1_min[f"{var}{ptype}_max"] = events_hem1_min[ f"EVT_thrutshemis0_{var}{ptype}"] #Recombine events = ak.concatenate([events_hem0_min, events_hem1_min], axis=0) #Plot the charged and neutral energies and multiplicities, where the low and high energy hemispheres are shown separately #Axis ranges, titles, and bins plot_config = { "echarged": [0, 50, "Charged energy [GeV]", 50], "eneutral": [0, 40, "Neutral energy [GeV]", 50], "ncharged": [0, 25, "Charged multiplicity", 25], "nneutral": [0, 20, "Neutral multiplicity", 20] } for p in plot_config: #Store the values in the tau container, so we can write it out tau[f"{p}_min"] = events[f"{p}_min"] tau[f"{p}_max"] = events[f"{p}_max"] fig, ax = plt.subplots(figsize=(9, 9)) plt.hist(events[f"{p}_min"], bins=plot_config[p][3], range=(plot_config[p][0], plot_config[p][1]), color="crimson", histtype='step', linewidth=2, label="Lower energy hemisphere per-event") plt.hist(events[f"{p}_max"], bins=plot_config[p][3], range=(plot_config[p][0], plot_config[p][1]), color="dodgerblue", histtype='step', linewidth=2, label="Higher energy hemisphere per-event") ax.tick_params(axis='both', which='major', labelsize=25) plt.xlim(plot_config[p][0], plot_config[p][1]) plt.xlabel(plot_config[p][2], fontsize=30) ymin, ymax = plt.ylim() plt.ylim(0, 1.15 * ymax) plt.legend(fontsize=18, loc="upper left") fig.savefig(f"{loc.PLOTS}/{B_type}2TauNu_min_max_hemisphere_{p}.pdf") #Plot the separation of the tau pion vertex from the PV (a measure of the B + tau flight) fig, ax = plt.subplots(figsize=(9, 9)) plt.hist(ak.flatten(tau["pi_1", "PVsep"]), range=(0, 10), bins=30, color="k", histtype='step', linewidth=2) ax.tick_params(axis='both', which='major', labelsize=25) plt.xlim(0, 10) plt.xlabel("$3\pi$ separation from PV [mm]", fontsize=30) fig.savefig(f"{loc.PLOTS}/{B_type}2TauNu_3pi_vertex_PV_sep.pdf") #Plot the separation of the true Bc decay vertex from the PV - this quantity cannot be reconstructed experimentally fig, ax = plt.subplots(figsize=(9, 9)) plt.hist(ak.flatten(tau["pi_1", "PVsep_parent"]), range=(0, 4), bins=30, color="k", histtype='step', linewidth=2) ax.tick_params(axis='both', which='major', labelsize=25) plt.xlim(0, 4) B_name = {"Bu": "$B^\\pm$", "Bc": "$B_c^\\pm$"} plt.xlabel(f"{B_name[B_type]} separation from PV [mm]", fontsize=30) fig.savefig(f"{loc.PLOTS}/{B_type}2TauNu_Bc_vertex_PV_sep.pdf") #Persist information into a flat dataframe for use elsewhere e.g. MVA training vs. inclusive Z -> qq/cc/bb background df = pd.DataFrame() persist_vars = [ "EVT_e", #Total event visible energy "EVT_thrutshemis_e_min", #Lowest energy hemisphere in event "EVT_thrutshemis_e_max", #Highest energy hemisphere in event "EVT_thrutshemis0_e", #costheta < 0 hemisphere energy "EVT_thrutshemis1_e", #costheta > 0 hemisphere energy "nominal_B_e", #Nominal signal B energy calculated using reco energy of true tau -> 3pi "px", "py", "pz", "e", "mass", "PVsep", #Distance of true tau decay vertex from the PV "PVsep_x", "PVsep_y", "PVsep_z", "PVsep_parent", #Distance of true Bc vertex from the PV (Bc flight which is not reconstructible) "EVT_thrust_x", "EVT_thrust_y", "EVT_thrust_z" ] for var in ["e", "n"]: for ptype in ["charged", "neutral"]: for m in ["min", "max"]: persist_vars.append(f"{var}{ptype}_{m}") for var in persist_vars: df[var] = ak.flatten(tau[var]).tolist() #Save to CSV ut.create_dir(loc.CSV) df.to_csv(f"{loc.CSV}/{B_type}2TauNu.csv") '''
def best_match(gen_hyp=None, jets=None, leptons=None, met=None): if gen_hyp is None: raise ValueError("Gen Objects gen_hyp needed for matching") if jets is None: raise ValueError("Reco jets needed for matching") if leptons is None: raise ValueError("Reco leptons needed for matching") if met is None: raise ValueError("Reco met needed for matching") if not ak.all(ak.num(gen_hyp) == 1): raise ValueError("Not all events for matching are semileptonic") jets_ak = ak.with_name(jets[["pt", "eta", "phi", "mass"]],"PtEtaPhiMLorentzVector") leps_ak = ak.with_name(leptons[["pt", "eta", "phi", "mass"]],"PtEtaPhiMLorentzVector") # init dict of objects matched_objects = {} # match jet closest to gen objects for genobj in ['BHad', 'BLep', 'WJa', 'WJb']: genobj_ak = ak.with_name(gen_hyp[genobj][["pt", "eta", "phi", "mass"]],"PtEtaPhiMLorentzVector") jets_akc, genobj_akc = ak.unzip(ak.cartesian([jets_ak, genobj_ak], nested=False)) deltaRs = jets_akc.delta_r(genobj_akc) # find deltaRs between jets and gen object indexOfMin = ak.unflatten(ak.argmin(deltaRs, axis=1), ak.num(genobj_ak)) passing_inds = deltaRs[indexOfMin] < 0.4 matched_jets_inds = indexOfMin[passing_inds] matched_jets = jets[matched_jets_inds] ## add matched perm objects matched_objects[genobj] = ak.Array({ 'pt' : matched_jets.pt, 'eta' : matched_jets.eta, 'phi' : matched_jets.phi, 'mass' : matched_jets.mass, 'jetIdx' : matched_jets_inds, # index of jet that the gen object is matched to in the event }, with_name="PtEtaPhiMLorentzVector") # match lepton closest to gen lepton genlep_ak = ak.with_name(gen_hyp['Lepton'][["pt", "eta", "phi", "mass"]],"PtEtaPhiMLorentzVector") lep_akc, genlep_akc = ak.unzip(ak.cartesian([leps_ak, genlep_ak], nested=False)) lepDRs = lep_akc.delta_r(genlep_akc) lepIdxOfMin = ak.unflatten(ak.argmin(lepDRs, axis=1), ak.num(genlep_ak)) passing_inds = lepDRs[lepIdxOfMin] < 0.4 matched_leps_inds = lepIdxOfMin[passing_inds] matched_leps = leptons[matched_leps_inds] ## add matched perm objects matched_objects['Lepton'] = ak.Array({key: matched_leps[key] for key in matched_leps.fields}, with_name="PtEtaPhiMLorentzVector") # solve for neutrino nu_array = np.zeros((len(ak.num(jets)), 4), dtype='float64') # convert all inputs into 2d numpy arrays of dtype=float64 (won't work if they're not float64) blep_inputs = np.stack((ak.to_numpy(ak.flatten(ak.fill_none(ak.pad_none(matched_objects['BLep'].px, 1), -999))).astype('float64'), ak.to_numpy(ak.flatten(ak.fill_none(ak.pad_none(matched_objects['BLep'].py, 1), -999))).astype('float64'),\ ak.to_numpy(ak.flatten(ak.fill_none(ak.pad_none(matched_objects['BLep'].pz, 1), -999))).astype('float64'), ak.to_numpy(ak.flatten(ak.fill_none(ak.pad_none(matched_objects['BLep'].energy, 1), -999))).astype('float64')), axis=-1) lep_inputs = np.stack((ak.to_numpy(ak.flatten(ak.fill_none(ak.pad_none(matched_objects['Lepton'].px, 1), -999))).astype('float64'), ak.to_numpy(ak.flatten(ak.fill_none(ak.pad_none(matched_objects['Lepton'].py, 1), -999))).astype('float64'),\ ak.to_numpy(ak.flatten(ak.fill_none(ak.pad_none(matched_objects['Lepton'].pz, 1), -999))).astype('float64'), ak.to_numpy(ak.flatten(ak.fill_none(ak.pad_none(matched_objects['Lepton'].energy, 1), -999))).astype('float64')), axis=-1) met_inputs = np.stack((ak.to_numpy(ak.fill_none(met.px, -999)).astype('float64'), ak.to_numpy(ak.fill_none(met.py, -999)).astype('float64')), axis=-1) nu_array = find_nu(bleps=blep_inputs, leptons=lep_inputs, met=met_inputs, nu_array=nu_array) valid_nu = ~((nu_array[:, 3] > 1e20) | (nu_array[:, 3] == 0)) # events that have a solution and matched blep # convert px, py, pz to pt, eta, phi nu_px, nu_py, nu_pz = nu_array[:, 0][valid_nu], nu_array[:, 1][valid_nu], nu_array[:, 2][valid_nu] nu_mom, nu_pt = np.sqrt(np.square(nu_px)+np.square(nu_py)+np.square(nu_pz)), np.sqrt(np.square(nu_px)+np.square(nu_py)) nu_phi = np.arctan2(nu_py, nu_px) nu_eta = np.arcsinh(nu_pz/nu_pt) matched_objects['Nu'] = ak.Array({ 'pt' : ak.unflatten(nu_pt, valid_nu.astype(int)), 'eta' : ak.unflatten(nu_eta, valid_nu.astype(int)), 'phi' : ak.unflatten(nu_phi, valid_nu.astype(int)), 'mass' : ak.zeros_like(ak.unflatten(nu_array[:, 0][valid_nu], valid_nu.astype(int))), 'chi2' : ak.unflatten(nu_array[:, 3][valid_nu], valid_nu.astype(int)), }, with_name="PtEtaPhiMLorentzVector") matched_perm = make_perm_table(bhad=matched_objects['BHad'], blep=matched_objects['BLep'], wja=matched_objects['WJa'], wjb=matched_objects['WJb'], lepton=matched_objects['Lepton'], met=met, nu=matched_objects['Nu']) return matched_perm
def get_evt_p4(p4, num=6): combos = ak.combinations(p4, num) part0, part1, part2, part3, part4, part5 = ak.unzip(combos) evt_p4 = part0 + part1 + part2 + part3 + part4 + part5 return evt_p4