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