def calculate_selection(self, syst_tag, events): """ """ electrons = events.ele electrons["label"] = -1 * awkward.ones_like(electrons.pt) if not self.is_data: electrons["label"] = awkward.where( electrons.genPartFlav == 1, awkward.ones_like(electrons.label), electrons.label) electrons["label"] = awkward.where( (electrons.genPartFlav == 3) | (electrons.genPartFlav == 4) | (electrons.genPartFlav == 5), awkward.zeros_like(electrons.label), electrons.label) fields = [x for x in events.fields if x not in ["ele", "Electron"]] for x in fields: if x == "ZCand": electrons[x] = awkward.firsts(events[x]) else: electrons[x] = events[x] electrons = awkward.flatten(electrons) dummy_cut = electrons.pt >= 0 return dummy_cut, electrons
def boost(self, other): """Apply a Lorentz boost given by the `ThreeVector` `other` and return it Note that this follows the convention that, for example in order to boost a vector into its own rest frame, one needs to use the negative of its `boostvec` """ b2 = other.rho2 gamma = (1 - b2)**(-0.5) mask = b2 == 0 b2 = awkward.where(mask, 1, b2) gamma2 = awkward.where(mask, 0, (gamma - 1) / b2) bp = self.dot(other) t = self.t v = gamma2 * bp * other + t * gamma * other return awkward.zip( { "x": self.x + v.x, "y": self.y + v.y, "z": self.z + v.z, "t": gamma * (t + bp), }, with_name="LorentzVector", behavior=self.behavior, )
def _kExtra(self, kpt, eta, nl, u, s=0, m=0): # if it is a jagged array, save the offsets then flatten everything # needed for the ternary conditions later abseta = abs(eta) kData = self._kRes[s][m][1](abseta) # type 1 is data kMC = self._kRes[s][m][0](abseta) # type 0 is MC mask = kData > kMC x = awkward.zeros_like(kpt) sigma = self._sigma(kpt, eta, nl, s, m) # Rochester cbA = beta, cbN = m, as well as cbM (always 0?) = loc and cbS = scale to transform y = (x-loc)/scale in the pdf method cbA = self._cbA[s][m](abseta, nl) cbN = self._cbN[s][m](abseta, nl) cbS = self._cbS[s][m](abseta, nl) counts = awkward.num(u) u_flat = awkward.flatten(u) loc = awkward.zeros_like(u_flat) cbA_flat = awkward.flatten(cbA) cbN_flat = awkward.flatten(cbN) cbS_flat = awkward.flatten(cbS) invcdf = awkward.unflatten( doublecrystalball.ppf(u_flat, cbA_flat, cbA_flat, cbN_flat, cbN_flat, loc, cbS_flat), counts, ) x = awkward.where( mask, (numpy.sqrt(kData * kData - kMC * kMC) * sigma * invcdf), x, ) result = awkward.where(x > -1, 1.0 / (1.0 + x), awkward.ones_like(kpt)) if isinstance(kpt, numpy.ndarray): result = numpy.array(result) return result
def rapidity(p_z): e_p = np.power(np.add(PROTON_MASS**2, np.power(p_z, 2)), 1/2) e_m = np.subtract(PROTON_MASS**2, np.power(p_z, 2)) e_m = ak.where(e_m < 0.0, 0.0, e_m) # to avoid imaginary numbers e_m = np.power(e_m, 1/2) e_m = ak.where(e_m == 0.0, 1e-10, e_m) # to avoid infinities y = np.multiply(np.log(np.divide(e_p, e_m)), 1/2) return y
def deltaPhi(jetphiL, metPhiL): phi1 = phi_(np.cos(jetphiL), np.sin(jetphiL)) phi2 = phi_(np.cos(metPhiL), np.sin(metPhiL)) dphi = phi1 - phi2 dphi_edited = ak.where(dphi < -np.pi, dphi + 2 * np.pi, dphi) dphi_edited = ak.where(dphi_edited > np.pi, dphi_edited - 2 * np.pi, dphi_edited) return abs(dphi_edited)
def deep_thr(tau, par, Pt_thr): c = 0.125 m = (c - par[0]) / (300 - Pt_thr) q = c - m * 300 thr1 = ak.where((tau.pt < 300), m * tau.pt + q, 0) thr2 = ak.where(tau.pt >= 300, c, 0) deep_thr = thr1 + thr2 return deep_thr
def test_issue_334(): a = ak.Array([1, 2, 3, 4]) b = ak.Array([-1]) c = ak.Array([True, False, True, True]) assert ak.where(c, a, b).tolist() == [1, -1, 3, 4] assert ak.where(*ak.broadcast_arrays(c, a, b)).tolist() == [1, -1, 3, 4] assert ak.where(c, a, -1).tolist() == [1, -1, 3, 4] assert ak.where(*ak.broadcast_arrays(c, a, -1)).tolist() == [1, -1, 3, 4]
def jer_smear(arr, variation, forceStochastic, ptGenJet, ptJet, etaJet): pt_gen = arr[ptGenJet] if not forceStochastic else None jetPt = arr[ptJet] if isinstance(jetPt, awkward.highlevel.Array): def getfunction(layout, depth): if (isinstance(layout, awkward.layout.NumpyArray) or not isinstance(layout, (awkward.layout.Content, awkward.partition.PartitionedArray))): return lambda: awkward.layout.NumpyArray( numpy.zeros_like(size=len(jetPt), dtype=numpy.float32)) return None if forceStochastic: pt_gen = awkward._util.recursively_apply( awkward.operations.convert.to_layout(jetPt), getfunction) pt_gen = awkward._util.wrap(pt_gen, awkward._util.behaviorof(jetPt)) else: raise Exception("'arr' must be an awkward array of some kind!") jersmear = arr['jet_energy_resolution'] * arr['jet_resolution_rand_gauss'] jersf = arr['jet_energy_resolution_scale_factor'][:, variation] doHybrid = pt_gen > 0 detSmear = 1 + (jersf - 1) * (jetPt - pt_gen) / jetPt # because of awkward1.0#367 stochSmear = 1 + numpy.sqrt(numpy.maximum(jersf**2 - 1, 0)) * jersmear min_jet_pt = _MIN_JET_ENERGY / numpy.cosh(arr[etaJet]) min_jet_pt_corr = min_jet_pt / jetPt smearfact = None if isinstance(arr, awkward.highlevel.Array): smearfact = awkward.where(doHybrid, detSmear, stochSmear) smearfact = awkward.where((smearfact * jetPt) < min_jet_pt, min_jet_pt_corr, smearfact) def getfunction(layout, depth): if (isinstance(layout, awkward.layout.NumpyArray) or not isinstance(layout, (awkward.layout.Content, awkward.partition.PartitionedArray))): return lambda: awkward.layout.NumpyArray(smearfact) return None smearfact = awkward._util.recursively_apply( awkward.operations.convert.to_layout(jetPt), getfunction) smearfact = awkward._util.wrap(smearfact, awkward._util.behaviorof(jetPt)) else: raise Exception("'arr' must be an awkward array of some kind!") return smearfact
def where(unique_keys): target_name = _hash_to_target_name[unique_keys[0]] mask = key == unique_keys[0] global_index = _get_global_index(events[target_name], obj._eventindex, index) global_index = awkward.where(mask, global_index, -1) links = events[target_name]._apply_global_index(global_index) if len(unique_keys) == 1: return links return awkward.where(mask, links, where(unique_keys[1:]))
def boostvec(self): """The `x`, `y` and `z` compontents divided by `t` as a `ThreeVector` This can be used for boosting. For cases where `|t| <= rho`, this returns the unit vector. """ rho = self.rho t = self.t with numpy.errstate(divide="ignore"): out = self.pvec * awkward.where( rho == 0, 0, awkward.where(abs(t) <= rho, 1 / rho, 1 / t)) return out
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 test_where(): condition = ak.Array([True, False, True, False, True], check_valid=True) one = ak.Array([1.1, 2.2, 3.3, 4.4, 5.5], check_valid=True) two = ak.Array([False, False, False, True, True], check_valid=True) three = ak.Array([[], [1], [2, 2], [3, 3, 3], [4, 4, 4, 4]], check_valid=True) assert ak.to_list(ak.where(condition, one, two)) == [1.1, 0.0, 3.3, 1.0, 5.5] assert ak.to_list(ak.where(condition, one, three)) == [ 1.1, [1], 3.3, [3, 3, 3], 5.5, ]
def Phi_mpi_pi(array_phi): """ Periodically (T=2*pi, shift step done only once) bring values of the given array into [-pi, pi] range. Arguments: array_phi: awkward array, values assumed to be radian measure of phi angle Returns: Awkward array, values of input array brought to [-pi, pi] range """ array_phi = ak.where(array_phi <= np.pi, array_phi, array_phi - 2*np.pi) array_phi = ak.where(array_phi >= -np.pi, array_phi, array_phi + 2*np.pi) assert ak.sum(array_phi > np.pi) + ak.sum(array_phi < -np.pi) == 0 return array_phi
def test_where(): one = ak.Array([0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9], check_valid=True) two = ak.Array([0, 100, 200, 300, 400, 500, 600, 700, 800, 900], check_valid=True) condition = ak.Array( [False, False, False, False, False, True, False, True, False, True], check_valid=True, ) assert isinstance(ak.where(condition)[0], ak.Array) assert ak.to_list(ak.where(condition)[0]) == [5, 7, 9] assert ak.to_list(ak.where(condition, one, two)) == ak.to_list( np.where(np.asarray(condition), np.asarray(one), np.asarray(two)))
def recalculate_flightDistance_from_arrays(value_branch, condition_branch): vertex_mask = (condition_branch == 1) | (condition_branch == 2) branch = value_branch.copy() zeros = ak.broadcast_arrays(0., branch)[0] ret_arr = ak.where(vertex_mask, zeros, branch) return ak.to_akward0(ret_arr)
def generate_epd_hit_matrix(self): ring_sum = np.zeros((32, len(self.nMip))) print("Filling array of dimension", ring_sum.shape) for i in range(32): ring_i = ak.sum(ak.where(self.row == i+1, self.nMip, 0), axis=-1) ring_sum[i] = ring_i return ring_sum
def _assignTruth(self, tree): nonSplitTrackSimClusIdx = self._getMatchIdxs(tree) truthEnergy = self._assignTruthByIndexAndSplit(tree,"MergedSimCluster_boundaryEnergy",nonSplitTrackSimClusIdx) truthPID = self._assignTruthByIndexAndSplit(tree,"MergedSimCluster_pdgId",nonSplitTrackSimClusIdx) truthX = self._assignTruthByIndexAndSplit(tree,"MergedSimCluster_impactPoint_x",nonSplitTrackSimClusIdx) truthY = self._assignTruthByIndexAndSplit(tree,"MergedSimCluster_impactPoint_y",nonSplitTrackSimClusIdx) truthZ = self._assignTruthByIndexAndSplit(tree,"MergedSimCluster_impactPoint_z",nonSplitTrackSimClusIdx) truthTime = self._assignTruthByIndexAndSplit(tree,"MergedSimCluster_impactPoint_t",nonSplitTrackSimClusIdx) #some manual sets zeros = ak1.zeros_like(truthEnergy) splittruthidx = self._splitJaggedArray(nonSplitTrackSimClusIdx) spectator = ak1.where(splittruthidx<0, zeros+10., zeros) trackPt = self._readSplitAndExpand(tree,"Track_pt") trackVertEta = self._readSplitAndExpand(tree,"Track_eta") trackMom = trackPt * np.cosh(trackVertEta) impactX = self._readSplitAndExpand(tree,"Track_HGCFront_x") impactY = self._readSplitAndExpand(tree,"Track_HGCFront_y") impactZ = self._readSplitAndExpand(tree,"Track_HGCFront_z") truthX = ak1.where(splittruthidx<0, impactX, truthX) truthY = ak1.where(splittruthidx<0, impactY, truthY) truthZ = ak1.where(splittruthidx<0, impactZ, truthZ) truthEnergy = ak1.where(splittruthidx<0, trackMom, truthEnergy) truthidx = self._expand(splittruthidx) self.truth={} self.truth['t_idx'] = truthidx# for now self.truth['t_energy'] = truthEnergy self.truth['t_pos'] = ak1.concatenate([truthX,truthY,truthZ],axis=-1) self.truth['t_time'] = truthTime self.truth['t_pid'] = truthPID self.truth['t_spectator'] = spectator self.truth['t_fully_contained'] = zeros+1 self.truth['t_rec_energy'] = trackMom
def _assignTruthByIndexAndSplit(self, tree, label, indices, null=0): sc = label if type(label) is str: sc = ak1.from_awkward0(tree[label].array()) vals = sc[indices] vals = ak1.where(indices<0, ak1.zeros_like(vals)+null, vals) ja = self._splitJaggedArray(vals) return self._expand(ja)
def jer_smear( variation, forceStochastic, pt_gen, jetPt, etaJet, jet_energy_resolution, jet_resolution_rand_gauss, jet_energy_resolution_scale_factor, ): pt_gen = pt_gen if not forceStochastic else None if not isinstance(jetPt, awkward.highlevel.Array): raise Exception("'jetPt' must be an awkward array of some kind!") if forceStochastic: pt_gen = awkward.without_parameters(awkward.zeros_like(jetPt)) jersmear = jet_energy_resolution * jet_resolution_rand_gauss jersf = jet_energy_resolution_scale_factor[:, variation] deltaPtRel = (jetPt - pt_gen) / jetPt doHybrid = (pt_gen > 0) & (numpy.abs(deltaPtRel) < 3 * jet_energy_resolution) detSmear = 1 + (jersf - 1) * deltaPtRel stochSmear = 1 + numpy.sqrt(numpy.maximum(jersf**2 - 1, 0)) * jersmear min_jet_pt = _MIN_JET_ENERGY / numpy.cosh(etaJet) min_jet_pt_corr = min_jet_pt / jetPt smearfact = awkward.where(doHybrid, detSmear, stochSmear) smearfact = awkward.where((smearfact * jetPt) < min_jet_pt, min_jet_pt_corr, smearfact) def getfunction(layout, depth): if isinstance(layout, awkward.layout.NumpyArray) or not isinstance( layout, (awkward.layout.Content, awkward.partition.PartitionedArray)): return lambda: awkward.layout.NumpyArray(smearfact) return None smearfact = awkward._util.recursively_apply( awkward.operations.convert.to_layout(jetPt), getfunction) smearfact = awkward._util.wrap(smearfact, awkward._util.behaviorof(jetPt)) return smearfact
def fsr_recovery(df): mask = ( (df.Muon.fsrPhotonIdx >= 0) & (df.Muon.matched_fsrPhoton.relIso03 < 1.8) & (df.Muon.matched_fsrPhoton.dROverEt2 < 0.012) & (df.Muon.matched_fsrPhoton.pt / df.Muon.pt < 0.4) & (abs(df.Muon.matched_fsrPhoton.eta) < 2.4) ) mask = ak.fill_none(mask, False) px = ak.zeros_like(df.Muon.pt) py = ak.zeros_like(df.Muon.pt) pz = ak.zeros_like(df.Muon.pt) e = ak.zeros_like(df.Muon.pt) fsr = { "pt": df.Muon.matched_fsrPhoton.pt, "eta": df.Muon.matched_fsrPhoton.eta, "phi": df.Muon.matched_fsrPhoton.phi, "mass": 0.0, } for obj in [df.Muon, fsr]: px_ = obj["pt"] * np.cos(obj["phi"]) py_ = obj["pt"] * np.sin(obj["phi"]) pz_ = obj["pt"] * np.sinh(obj["eta"]) e_ = np.sqrt(px_**2 + py_**2 + pz_**2 + obj["mass"] ** 2) px = px + px_ py = py + py_ pz = pz + pz_ e = e + e_ pt = np.sqrt(px**2 + py**2) eta = np.arcsinh(pz / pt) phi = np.arctan2(py, px) mass = np.sqrt(e**2 - px**2 - py**2 - pz**2) iso = (df.Muon.pfRelIso04_all * df.Muon.pt - df.Muon.matched_fsrPhoton.pt) / pt df["Muon", "pt_fsr"] = ak.where(mask, pt, df.Muon.pt) df["Muon", "eta_fsr"] = ak.where(mask, eta, df.Muon.eta) df["Muon", "phi_fsr"] = ak.where(mask, phi, df.Muon.phi) df["Muon", "mass_fsr"] = ak.where(mask, mass, df.Muon.mass) df["Muon", "iso_fsr"] = ak.where(mask, iso, df.Muon.pfRelIso04_all) return mask
def recalculate_flightDistance(tree, key): vertex_mask = (tree["TagVarCSV_vertexCategory"].array() == 1) | (tree["TagVarCSV_vertexCategory"].array() == 2) branch = tree[key].array().copy() zeros = ak.broadcast_arrays(0., branch)[0] ret_arr = ak.where(vertex_mask, zeros, branch) return ak.to_awkward0(ret_arr)
def ak_solve_4PJ(self, mthad, mwhad, nschi): "Inputs: m(thad_proxy), m(whad_proxy), and the NS solver sqrt(chi2) values (mthad, mwhad, nschi)\nOutputs: numpy array of the mass and NS disriminants (MassDiscr, NuDiscr)" #set_trace() mass_discval = ak.where((mthad < self.WTmass_right._axes[0][-1]) & (mwhad < self.WTmass_right._axes[1][-1]), self.WTmass_right(mthad, mwhad) / np.sum(self.WTmass_right._values), ak.ones_like(mthad) * np.nan) masstest = -1 * np.log(mass_discval) masstest = ak.nan_to_num(masstest, nan=np.inf) nu_discval = ak.where( nschi < self.NS_4PJ_right._axes[-1], self.NS_4PJ_right(nschi) / np.sum(self.NS_4PJ_right._values), ak.ones_like(nschi) * np.nan) nstest = -1 * np.log(nu_discval) nstest = ak.nan_to_num(nstest, nan=np.inf) return masstest, nstest
def __init__(self, mID, mQT_data, mnMips, lower_bound=0.2, upper_bound=3): self.mID = mID self.mQT_data = mQT_data self.mnMip = mnMips self.has_TAC = np.bitwise_and(np.right_shift(self.mQT_data, 29), 0x1) self.status_is_good = np.bitwise_and(np.right_shift(self.mQT_data, 30), 0x1) self.adc = np.bitwise_and(self.mQT_data, 0x0FFF) self.tac = np.bitwise_and(np.right_shift(self.mQT_data, 12), 0x0FFF) self.TDC = np.bitwise_and(np.right_shift(self.mQT_data, 24), 0x001F) self.EW = np.sign(self.mID) self.position = np.abs(self.mID // 100) self.tile = np.abs(self.mID) % 100 self.row = np.abs(mID) % 100 // 2 + 1 self.nMip = ak.where(self.status_is_good, self.mnMip, 0) # nMIP truncation self.nMip = ak.where(self.nMip <= lower_bound, lower_bound, self.nMip) self.nMip = ak.where(self.nMip >= upper_bound, upper_bound, self.nMip)
def ak_solve_3J_lost(self, m2jets, nschi): "Inputs: m(j1+j2) and the NS solver sqrt(chi2) values (m2jets, nschi)\nOutputs: numpy array of the combined, mass, and NS disriminants (MassDiscr, NuDiscr)" #set_trace() mass_discval = ak.where( m2jets < self.Mass_3J_Lost_right._axes[-1], self.Mass_3J_Lost_right(m2jets) / np.sum(self.Mass_3J_Lost_right._values), ak.ones_like(m2jets) * np.nan) masstest = -1 * np.log(mass_discval) masstest = ak.nan_to_num(masstest, nan=np.inf) nu_discval = ak.where( nschi < self.NS_3J_Lost_right._axes[-1], self.NS_3J_Lost_right(nschi) / np.sum(self.NS_3J_Lost_right._values), ak.ones_like(nschi) * np.nan) nstest = -1 * np.log(nu_discval) nstest = ak.nan_to_num(nstest, nan=np.inf) return masstest, nstest
def test_getitem_field(): a1 = ak.zip({ "a": [[1], [], [2, 3]], "b": [[4], [], [5, 6]] }, with_name="a1") a2 = ak.zip({ "a": [[7, 8], [9], []], "b": [[10, 11], [12], []] }, with_name="a2") union = ak.where([True, False, True], a1, a2) assert str(union.a.type) == "3 * var * int64"
def test_flatten_axis_none(): a1 = ak.zip({ "a": [[1], [], [2, 3]], "b": [[4], [], [5, 6]] }, with_name="a1") a2 = ak.zip({ "a": [[7, 8], [9], []], "b": [[10, 11], [12], []] }, with_name="a2") union = ak.where([True, False, True], a1, a2) assert set(ak.flatten(union, axis=None)) == {1, 2, 3, 4, 5, 6, 9, 12}
def test(): one = ak.Array([[0, 1, 2], [], [3, 4], [5], [6, 7, 8, 9]]) two = ak.Array([[0.0, 1.1, 2.2], [], [3.3, 4.4], [5.5], [6.6, 7.7, 8.8, 9.9]]) condition = ak.Array([[False, True, False], [], [True, False], [True], [False, False, True, True]]) assert ak.where(condition, one, two).tolist() == [ [0, 1, 2.2], [], [3, 4.4], [5], [6.6, 7.7, 8, 9], ]
def test(): array = ak.Array( ak.layout.RegularArray( ak.layout.NumpyArray(np.r_[1, 2, 3, 4, 5, 6, 7, 8, 9]), 3)) condition = ak.Array( ak.layout.NumpyArray( np.array([[True, True, True], [True, True, False], [True, False, True]]))) assert ak.where(condition == 2, array, 2 * array).tolist() == [ [2, 4, 6], [8, 10, 12], [14, 16, 18], ]
def apply_geofit(df, year, mask): d0_BS_charge = np.multiply(df.Muon.dxybs, df.Muon.charge) mask = mask & (np.abs(d0_BS_charge) < 999999.0) pt = df.Muon.pt eta = df.Muon.eta cuts = { "eta_1": (np.abs(eta) < 0.9), "eta_2": ((np.abs(eta) < 1.7) & (np.abs(eta) >= 0.9)), "eta_3": (np.abs(eta) >= 1.7), } factors = { "2016": { "eta_1": 411.34, "eta_2": 673.40, "eta_3": 1099.0 }, "2017": { "eta_1": 582.32, "eta_2": 974.05, "eta_3": 1263.4 }, "2018": { "eta_1": 650.84, "eta_2": 988.37, "eta_3": 1484.6 }, } pt_corr = pt for eta_i in ["eta_1", "eta_2", "eta_3"]: value = factors[year][eta_i] * d0_BS_charge * pt * pt / 10000.0 pt_corr = ak.where(cuts[eta_i], value, pt_corr) df["Muon", "pt_gf"] = ak.where(mask, pt - pt_corr, pt) return
def get_vpt(check_offshell=False): """Only the leptonic samples have no resonance in the decay tree, and only when M is beyond the configured Breit-Wigner cutoff (usually 15*width) """ boson = ak.firsts( genpart[((genpart.pdgId == 23) | (abs(genpart.pdgId) == 24)) & genpart.hasFlags(["fromHardProcess", "isLastCopy"])]) if check_offshell: offshell = genpart[ genpart.hasFlags(["fromHardProcess", "isLastCopy"]) & ak.is_none(boson) & (abs(genpart.pdgId) >= 11) & (abs(genpart.pdgId) <= 16)].sum() return ak.where(ak.is_none(boson.pt), offshell.pt, boson.pt) return np.array(ak.fill_none(boson.pt, 0.))