def test():
original = ak.Array([[1, 2, 3, 4], [], [5, 6, 7], [8, 9]])
assert ak.unflatten(original, [2, 2, 1, 2, 1, 1], axis=1).tolist() == [
[[1, 2], [3, 4]],
[],
[[5], [6, 7]],
[[8], [9]],
]
assert ak.unflatten(original, [1, 3, 1, 2, 1, 1], axis=1).tolist() == [
[[1], [2, 3, 4]],
[],
[[5], [6, 7]],
[[8], [9]],
]
with pytest.raises(ValueError):
ak.unflatten(original, [2, 1, 2, 2, 1, 1], axis=1)
assert ak.unflatten(original, [2, 0, 2, 1, 2, 1, 1], axis=1).tolist() == [
[[1, 2], [], [3, 4]],
[],
[[5], [6, 7]],
[[8], [9]],
]
def test():
array = ak.Array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
assert ak.unflatten(array, 5).tolist() == [[0, 1, 2, 3, 4],
[5, 6, 7, 8, 9]]
assert ak.unflatten(array, [3, 0, 2, 1, 4]).tolist() == [
[0, 1, 2],
[],
[3, 4],
[5],
[6, 7, 8, 9],
]
assert ak.unflatten(array, [3, None, 2, 1, 4]).tolist() == [
[0, 1, 2],
None,
[3, 4],
[5],
[6, 7, 8, 9],
]
original = ak.Array([[0, 1, 2], [], [3, 4], [5], [6, 7, 8, 9]])
counts = ak.num(original)
array = ak.flatten(original)
assert counts.tolist() == [3, 0, 2, 1, 4]
assert array.tolist() == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
assert ak.unflatten(array, counts).tolist() == [
[0, 1, 2],
[],
[3, 4],
[5],
[6, 7, 8, 9],
]
def test_jet_correction_uncertainty_sources():
from coffea.jetmet_tools import JetCorrectionUncertainty
counts, test_eta, test_pt = dummy_jagged_eta_pt()
test_pt_jag = ak.unflatten(test_pt, counts)
test_eta_jag = ak.unflatten(test_eta, counts)
junc_names = []
levels = []
for name in dir(evaluator):
if 'Summer16_23Sep2016V3_MC_UncertaintySources_AK4PFPuppi' in name:
junc_names.append(name)
levels.append(name.split('_')[-1])
#test for underscore in dataera
if 'Fall17_17Nov2017_V6_MC_UncertaintySources_AK4PFchs_AbsoluteFlavMap' in name:
junc_names.append(name)
levels.append(name.split('_')[-1])
junc = JetCorrectionUncertainty(
**{name: evaluator[name]
for name in junc_names})
print(junc)
juncs = junc.getUncertainty(JetEta=test_eta, JetPt=test_pt)
juncs_jag = list(
junc.getUncertainty(JetEta=test_eta_jag, JetPt=test_pt_jag))
for i, (level, corrs) in enumerate(juncs):
assert (level in levels)
assert (corrs.shape[0] == test_eta.shape[0])
tic = time.time()
assert (ak.all(corrs == ak.flatten(juncs_jag[i][1])))
toc = time.time()
def test_groupby():
array = ak.Array(
[
{"x": 1, "y": 1.1},
{"x": 2, "y": 2.2},
{"x": 1, "y": 1.1},
{"x": 3, "y": 3.3},
{"x": 1, "y": 1.1},
{"x": 2, "y": 2.2},
]
)
sorted = array[ak.argsort(array.x)]
assert sorted.x.tolist() == [1, 1, 1, 2, 2, 3]
assert ak.run_lengths(sorted.x).tolist() == [3, 2, 1]
assert ak.unflatten(sorted, ak.run_lengths(sorted.x)).tolist() == [
[{"x": 1, "y": 1.1}, {"x": 1, "y": 1.1}, {"x": 1, "y": 1.1}],
[{"x": 2, "y": 2.2}, {"x": 2, "y": 2.2}],
[{"x": 3, "y": 3.3}],
]
array = ak.Array(
[
[{"x": 1, "y": 1.1}, {"x": 2, "y": 2.2}, {"x": 1, "y": 1.1}],
[{"x": 3, "y": 3.3}, {"x": 1, "y": 1.1}, {"x": 2, "y": 2.2}],
]
)
sorted = array[ak.argsort(array.x)]
assert sorted.x.tolist() == [[1, 1, 2], [1, 2, 3]]
assert ak.run_lengths(sorted.x).tolist() == [[2, 1], [1, 1, 1]]
counts = ak.flatten(ak.run_lengths(sorted.x), axis=None)
assert ak.unflatten(sorted, counts, axis=-1).tolist() == [
[[{"x": 1, "y": 1.1}, {"x": 1, "y": 1.1}], [{"x": 2, "y": 2.2}]],
[[{"x": 1, "y": 1.1}], [{"x": 2, "y": 2.2}], [{"x": 3, "y": 3.3}]],
]
def test_jet_correction_regrouped_uncertainty_sources():
from coffea.jetmet_tools import JetCorrectionUncertainty
counts, test_eta, test_pt = dummy_jagged_eta_pt()
test_pt_jag = ak.unflatten(test_pt, counts)
test_eta_jag = ak.unflatten(test_eta, counts)
junc_names = []
levels = []
for name in dir(evaluator):
if 'Regrouped_Fall17_17Nov2017_V32_MC_UncertaintySources_AK4PFchs' in name:
junc_names.append(name)
if len(name.split('_')) == 9:
levels.append("_".join(name.split('_')[-2:]))
else:
levels.append(name.split('_')[-1])
junc = JetCorrectionUncertainty(
**{name: evaluator[name]
for name in junc_names})
print(junc)
juncs_jag = list(
junc.getUncertainty(JetEta=test_eta_jag, JetPt=test_pt_jag))
for i, tpl in enumerate(
list(junc.getUncertainty(JetEta=test_eta, JetPt=test_pt))):
assert (tpl[0] in levels)
assert (tpl[1].shape[0] == test_eta.shape[0])
assert (ak.all(tpl[1] == ak.flatten(juncs_jag[i][1])))
def test():
content = ak.repartition(range(10), 3)
assert ak.unflatten(content, [3, 0, 2, 1, 3, 1]).tolist() == [
[0, 1, 2],
[],
[3, 4],
[5],
[6, 7, 8],
[9],
]
original = ak.repartition(ak.Array([[1, 2, 3, 4], [5, 6, 7], [8, 9]]), 2)
assert ak.unflatten(original, [2, 2, 1, 2, 1, 1], axis=1).tolist() == [
[[1, 2], [3, 4]],
[[5], [6, 7]],
[[8], [9]],
]
original = ak.repartition(ak.Array([[1, 2, 3, 4], [], [5, 6, 7], [8, 9]]),
2)
assert ak.unflatten(original, [2, 2, 1, 2, 1, 1], axis=1).tolist() == [
[[1, 2], [3, 4]],
[],
[[5], [6, 7]],
[[8], [9]],
]
def test_root_scalefactors():
extractor = lookup_tools.extractor()
extractor.add_weight_sets([
"testSF2d scalefactors_Tight_Electron tests/samples/testSF2d.histo.root"
])
extractor.finalize(reduce_list=["testSF2d"])
evaluator = extractor.make_evaluator()
counts, test_eta, test_pt = dummy_jagged_eta_pt()
# test flat eval
test_out = evaluator["testSF2d"](test_eta, test_pt)
# print it
print(evaluator["testSF2d"])
# test structured eval
test_eta_jagged = ak.unflatten(test_eta, counts)
test_pt_jagged = ak.unflatten(test_pt, counts)
test_out_jagged = evaluator["testSF2d"](test_eta_jagged, test_pt_jagged)
assert ak.all(ak.num(test_out_jagged) == counts)
assert ak.all(ak.flatten(test_out_jagged) == test_out)
print(test_out)
diff = np.abs(test_out - _testSF2d_expected_output)
print("Max diff: %.16f" % diff.max())
print("Median diff: %.16f" % np.median(diff))
print("Diff over threshold rate: %.1f %%" %
(100 * (diff >= 1.0e-8).sum() / diff.size))
assert (diff < 1.0e-8).all()
def test():
a = ak.Array([[1, 2, 3, 4], [5, 6, 7, 8]])
assert ak.unflatten(a, [2, 2, 2, 2], axis=1).tolist() == [
[[1, 2], [3, 4]],
[[5, 6], [7, 8]],
]
assert ak.unflatten(a, 2, axis=1).tolist() == [[[1, 2], [3, 4]],
[[5, 6], [7, 8]]]
def test():
array = ak.Array([[1, 2, 3]])
assert ak.unflatten(array, [2, 1], axis=1).tolist() == [[[1, 2], [3]]]
assert ak.unflatten(array, [0, 2, 1], axis=1).tolist() == [[[], [1, 2],
[3]]]
assert ak.unflatten(array, [0, 0, 2, 1], axis=1).tolist() == [[[], [],
[1, 2],
[3]]]
def apply_roccor(df, rochester, is_mc):
if is_mc:
hasgen = ~np.isnan(ak.fill_none(df.Muon.matched_gen.pt, np.nan))
mc_rand = np.random.rand(*ak.to_numpy(ak.flatten(df.Muon.pt)).shape)
mc_rand = ak.unflatten(mc_rand, ak.num(df.Muon.pt, axis=1))
corrections = np.array(ak.flatten(ak.ones_like(df.Muon.pt)))
errors = np.array(ak.flatten(ak.ones_like(df.Muon.pt)))
mc_kspread = rochester.kSpreadMC(
df.Muon.charge[hasgen],
df.Muon.pt[hasgen],
df.Muon.eta[hasgen],
df.Muon.phi[hasgen],
df.Muon.matched_gen.pt[hasgen],
)
mc_ksmear = rochester.kSmearMC(
df.Muon.charge[~hasgen],
df.Muon.pt[~hasgen],
df.Muon.eta[~hasgen],
df.Muon.phi[~hasgen],
df.Muon.nTrackerLayers[~hasgen],
mc_rand[~hasgen],
)
errspread = rochester.kSpreadMCerror(
df.Muon.charge[hasgen],
df.Muon.pt[hasgen],
df.Muon.eta[hasgen],
df.Muon.phi[hasgen],
df.Muon.matched_gen.pt[hasgen],
)
errsmear = rochester.kSmearMCerror(
df.Muon.charge[~hasgen],
df.Muon.pt[~hasgen],
df.Muon.eta[~hasgen],
df.Muon.phi[~hasgen],
df.Muon.nTrackerLayers[~hasgen],
mc_rand[~hasgen],
)
hasgen_flat = np.array(ak.flatten(hasgen))
corrections[hasgen_flat] = np.array(ak.flatten(mc_kspread))
corrections[~hasgen_flat] = np.array(ak.flatten(mc_ksmear))
errors[hasgen_flat] = np.array(ak.flatten(errspread))
errors[~hasgen_flat] = np.array(ak.flatten(errsmear))
corrections = ak.unflatten(corrections, ak.num(df.Muon.pt, axis=1))
errors = ak.unflatten(errors, ak.num(df.Muon.pt, axis=1))
else:
corrections = rochester.kScaleDT(df.Muon.charge, df.Muon.pt,
df.Muon.eta, df.Muon.phi)
errors = rochester.kScaleDTerror(df.Muon.charge, df.Muon.pt,
df.Muon.eta, df.Muon.phi)
df["Muon", "pt_roch"] = df.Muon.pt * corrections
df["Muon", "pt_roch_up"] = df.Muon.pt_roch + df.Muon.pt * errors
df["Muon", "pt_roch_down"] = df.Muon.pt_roch - df.Muon.pt * errors
def test_jet_correction_uncertainty():
from coffea.jetmet_tools import JetCorrectionUncertainty
counts, test_eta, test_pt = dummy_jagged_eta_pt()
test_pt_jag = ak.unflatten(test_pt, counts)
test_eta_jag = ak.unflatten(test_eta, counts)
junc_names = ["Summer16_23Sep2016V3_MC_Uncertainty_AK4PFPuppi"]
junc = JetCorrectionUncertainty(
**{name: evaluator[name]
for name in junc_names})
print(junc)
juncs = junc.getUncertainty(JetEta=test_eta, JetPt=test_pt)
juncs_jag = list(
junc.getUncertainty(JetEta=test_eta_jag, JetPt=test_pt_jag))
for i, (level, corrs) in enumerate(juncs):
assert corrs.shape[0] == test_eta.shape[0]
assert ak.all(corrs == ak.flatten(juncs_jag[i][1]))
test_pt_jag = test_pt_jag[0:3]
test_eta_jag = test_eta_jag[0:3]
counts = counts[0:3]
print("Raw jet values:")
print("pT:", test_pt_jag.tolist())
print("eta:", test_eta_jag.tolist(), "\n")
juncs_jag_ref = ak.unflatten(
np.array([
[1.053504214, 0.946495786],
[1.033343349, 0.966656651],
[1.065159157, 0.934840843],
[1.033140127, 0.966859873],
[1.016858652, 0.983141348],
[1.130199999, 0.869800001],
[1.039968468, 0.960031532],
[1.033100002, 0.966899998],
]),
counts,
)
juncs_jag = list(
junc.getUncertainty(JetEta=test_eta_jag, JetPt=test_pt_jag))
for i, (level, corrs) in enumerate(juncs_jag):
print("Index:", i)
print("Correction level:", level)
print("Reference Uncertainties (jagged):", juncs_jag_ref)
print("Uncertainties (jagged):", corrs)
assert ak.all(
np.abs(ak.flatten(juncs_jag_ref) - ak.flatten(corrs)) < 1e-6)
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_jet_resolution_sf_2d():
from coffea.jetmet_tools import JetResolutionScaleFactor
counts, test_eta, test_pt = dummy_jagged_eta_pt()
test_pt_jag = ak.unflatten(test_pt, counts)
test_eta_jag = ak.unflatten(test_eta, counts)
resosf = JetResolutionScaleFactor(
**{name: evaluator[name]
for name in ["Autumn18_V7_MC_SF_AK4PFchs"]})
resosfs = resosf.getScaleFactor(JetPt=test_pt, JetEta=test_eta)
resosfs_jag = resosf.getScaleFactor(JetPt=test_pt_jag, JetEta=test_eta_jag)
def test_root_scalefactors():
extractor = lookup_tools.extractor()
extractor.add_weight_sets([
"testSF2d scalefactors_Tight_Electron tests/samples/testSF2d.histo.root"
])
extractor.finalize(reduce_list=['testSF2d'])
evaluator = extractor.make_evaluator()
counts, test_eta, test_pt = dummy_jagged_eta_pt()
# test flat eval
test_out = evaluator["testSF2d"](test_eta, test_pt)
# print it
print(evaluator["testSF2d"])
# test structured eval
test_eta_jagged = ak.unflatten(test_eta, counts)
test_pt_jagged = ak.unflatten(test_pt, counts)
test_out_jagged = evaluator["testSF2d"](test_eta_jagged, test_pt_jagged)
assert ak.all(ak.num(test_out_jagged) == counts)
assert ak.all(ak.flatten(test_out_jagged) == test_out)
# From make_expected_lookup.py
expected_output = np.array([
0.90780139, 0.82748538, 0.86332178, 0.86332178, 0.97981155, 0.79701495,
0.88245934, 0.82857144, 0.91884059, 0.97466666, 0.94072163, 1.00775194,
0.82748538, 1.00775194, 0.97203946, 0.98199672, 0.80655736, 0.90893763,
0.88245934, 0.79701495, 0.82748538, 0.82857144, 0.91884059, 0.90893763,
0.97520661, 0.97520661, 0.82748538, 0.91884059, 0.97203946, 0.88245934,
0.79701495, 0.9458763, 1.00775194, 0.80655736, 1.00775194, 1.00775194,
0.98976982, 0.98976982, 0.86332178, 0.94072163, 0.80655736, 0.98976982,
0.96638656, 0.9458763, 0.90893763, 0.9529984, 0.9458763, 0.9529984,
0.80655736, 0.80655736, 0.80655736, 0.98976982, 0.97466666, 0.98199672,
0.86332178, 1.03286386, 0.94072163, 1.03398061, 0.82857144, 0.80655736,
1.00775194, 0.80655736
])
print(test_out)
diff = np.abs(test_out - expected_output)
print("Max diff: %.16f" % diff.max())
print("Median diff: %.16f" % np.median(diff))
print("Diff over threshold rate: %.1f %%" %
(100 * (diff >= 1.e-8).sum() / diff.size))
assert (diff < 1.e-8).all()
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 load(buffer, entry_start: int = None, entry_stop: int = None) -> ak.Array:
"""Load MIDAS data from a buffer and return the ragged array.
:param buffer: object supporting buffer protocol
:param entry_start: the first entry to include. If None, start at 0.
:param entry_stop: the first entry to exclude. If None, stop at the end of the file.
"""
if entry_start is None:
entry_start = 0
if entry_stop is None:
entry_stop = np.iinfo(np.int32).max
if entry_start < 0:
raise ValueError("Start entry must be positive")
if entry_stop < 0:
raise ValueError("Stop entry must be positive")
if entry_stop < entry_start:
raise ValueError("Stop entry must not be less than the start entry")
# Create a view over the memory map
data = np.frombuffer(buffer, dtype="<B")
# Run the VM
vm = ForthMachine32(MIDAS_PROGRAM)
vm.run({
"data": data,
"range": np.array([entry_start, entry_stop], dtype=np.int32)
})
# Create the result
outputs = vm.outputs.copy()
num = outputs.pop("num")
return ak.unflatten(ak.zip(outputs), num)
def constituents(self, min_pt):
outputs_to_inputs = self.constituent_index(min_pt)
shape = ak.num(outputs_to_inputs)
total = np.sum(shape)
duplicate = ak.unflatten(np.zeros(total, np.int64), shape)
prepared = self.data[:, np.newaxis][duplicate]
return prepared[outputs_to_inputs]
def test_jet_resolution_sf():
from coffea.jetmet_tools import JetResolutionScaleFactor
counts, test_eta, test_pt = dummy_jagged_eta_pt()
test_pt_jag = ak.unflatten(test_pt, counts)
test_eta_jag = ak.unflatten(test_eta, counts)
jersf_names = ["Spring16_25nsV10_MC_SF_AK4PFPuppi"]
resosf = JetResolutionScaleFactor(
**{name: evaluator[name]
for name in jersf_names})
print(resosf)
# 0-jet compatibility
assert resosf.getScaleFactor(JetEta=test_eta[:0]).shape == (0, 3)
resosfs = resosf.getScaleFactor(JetEta=test_eta)
resosfs_jag = resosf.getScaleFactor(JetEta=test_eta_jag)
assert ak.all(resosfs == ak.flatten(resosfs_jag))
test_pt_jag = test_pt_jag[0:3]
test_eta_jag = test_eta_jag[0:3]
counts = counts[0:3]
print("Raw jet values:")
print("pT:", test_pt_jag)
print("eta:", test_eta_jag, "\n")
resosfs_jag_ref = ak.unflatten(
np.array([
[1.857, 1.928, 1.786],
[1.084, 1.095, 1.073],
[1.364, 1.403, 1.325],
[1.177, 1.218, 1.136],
[1.138, 1.151, 1.125],
[1.364, 1.403, 1.325],
[1.177, 1.218, 1.136],
[1.082, 1.117, 1.047],
]),
counts,
)
resosfs_jag = resosf.getScaleFactor(JetEta=test_eta_jag)
print("Reference Resolution SF (jagged):", resosfs_jag_ref)
print("Resolution SF (jagged):", resosfs_jag)
assert ak.all(
np.abs(ak.flatten(resosfs_jag_ref) - ak.flatten(resosfs_jag)) < 1e-6)
def test_jet_resolution_sf_2d():
from coffea.jetmet_tools import JetResolutionScaleFactor
counts, test_eta, test_pt = dummy_jagged_eta_pt()
test_pt_jag = ak.unflatten(test_pt, counts)
test_eta_jag = ak.unflatten(test_eta, counts)
resosf = JetResolutionScaleFactor(
**{name: evaluator[name]
for name in ["Autumn18_V7_MC_SF_AK4PFchs"]})
print(resosf)
# 0-jet compatibility
assert resosf.getScaleFactor(JetPt=test_pt[:0],
JetEta=test_eta[:0]).shape == (0, 3)
resosfs = resosf.getScaleFactor(JetPt=test_pt, JetEta=test_eta)
resosfs_jag = resosf.getScaleFactor(JetPt=test_pt_jag, JetEta=test_eta_jag)
assert ak.all(resosfs == ak.flatten(resosfs_jag))
test_pt_jag = test_pt_jag[0:3]
test_eta_jag = test_eta_jag[0:3]
counts = counts[0:3]
print("Raw jet values:")
print("pT:", test_pt_jag)
print("eta:", test_eta_jag, "\n")
resosfs_jag_ref = ak.unflatten(
np.array([
[1.11904, 1.31904, 1.0],
[1.1432, 1.2093, 1.0771],
[1.16633, 1.36633, 1.0],
[1.17642, 1.37642, 1.0],
[1.1808, 1.1977, 1.1640],
[1.15965, 1.35965, 1.0],
[1.17661, 1.37661, 1.0],
[1.1175, 1.1571, 1.0778],
]),
counts,
)
resosfs_jag = resosf.getScaleFactor(JetPt=test_pt_jag, JetEta=test_eta_jag)
print("Reference Resolution SF (jagged):", resosfs_jag_ref)
print("Resolution SF (jagged):", resosfs_jag)
assert ak.all(
np.abs(ak.flatten(resosfs_jag_ref) - ak.flatten(resosfs_jag)) < 1e-6)
def test_evaluator_exceptions():
extractor = lookup_tools.extractor()
extractor.add_weight_sets([
"testSF2d scalefactors_Tight_Electron tests/samples/testSF2d.histo.root"
])
counts, test_eta, test_pt = dummy_jagged_eta_pt()
test_eta_jagged = ak.unflatten(test_eta, counts)
test_pt_jagged = ak.unflatten(test_pt, counts)
extractor.finalize()
evaluator = extractor.make_evaluator()
try:
test_out = evaluator["testSF2d"](test_pt_jagged, test_eta)
except Exception as e:
assert (isinstance(e, TypeError))
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, 100.)
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 = reso.getResolution(JetEta=test_eta_jag,
Rho=test_Rho_jag,
JetPt=test_pt_jag)
def test_jet_resolution_sf():
from coffea.jetmet_tools import JetResolutionScaleFactor
counts, test_eta, test_pt = dummy_jagged_eta_pt()
test_pt_jag = ak.unflatten(test_pt, counts)
test_eta_jag = ak.unflatten(test_eta, counts)
jersf_names = ['Spring16_25nsV10_MC_SF_AK4PFPuppi']
resosf = JetResolutionScaleFactor(
**{name: evaluator[name]
for name in jersf_names})
print(resosf)
resosfs = resosf.getScaleFactor(JetEta=test_eta)
resosfs_jag = resosf.getScaleFactor(JetEta=test_eta_jag)
assert (ak.all(resosfs == ak.flatten(resosfs_jag)))
def test():
layout = ak.layout.IndexedArray64(
ak.layout.Index64(np.array([3, 1, 0, 2])),
ak.layout.ListOffsetArray64(
ak.layout.Index64(np.array([0, 3, 6, 9, 12])),
ak.layout.NumpyArray(np.array([0, 0, 0, 1, 1, 1, 2, 2, 3, 3, 3, 3])),
),
)
assert ak.unflatten(
layout, ak.flatten(ak.run_lengths(layout)), axis=1
).tolist() == [[[3, 3, 3]], [[1, 1, 1]], [[0, 0, 0]], [[2, 2], [3]]]
def test_factorized_jet_corrector():
from coffea.jetmet_tools import FactorizedJetCorrector
counts, test_eta, test_pt = dummy_jagged_eta_pt()
test_Rho = np.full_like(test_eta, 100.)
test_A = np.full_like(test_eta, 5.)
jec_names = [
'Summer16_23Sep2016V3_MC_L1FastJet_AK4PFPuppi',
'Summer16_23Sep2016V3_MC_L2Relative_AK4PFPuppi',
'Summer16_23Sep2016V3_MC_L2L3Residual_AK4PFPuppi',
'Summer16_23Sep2016V3_MC_L3Absolute_AK4PFPuppi'
]
corrector = FactorizedJetCorrector(
**{name: evaluator[name]
for name in jec_names})
print(corrector)
pt_copy = np.copy(test_pt)
corrs = corrector.getCorrection(JetEta=test_eta,
Rho=test_Rho,
JetPt=test_pt,
JetA=test_A)
assert ((np.abs(pt_copy - test_pt) < 1e-6).all())
test_pt_jag = ak.unflatten(test_pt, counts)
test_eta_jag = ak.unflatten(test_eta, counts)
test_Rho_jag = ak.unflatten(test_Rho, counts)
test_A_jag = ak.unflatten(test_A, counts)
corrs_jag = corrector.getCorrection(JetEta=test_eta_jag,
Rho=test_Rho_jag,
JetPt=test_pt_jag,
JetA=test_A_jag)
assert (ak.all(np.abs(pt_copy - ak.flatten(test_pt_jag)) < 1e-6))
assert (ak.all(np.abs(corrs - ak.flatten(corrs_jag)) < 1e-6))
def counts2nestedindex(stack):
"""Turn jagged local counts into doubly-kagged global index into a target
Signature: local_counts,target_offsets,!counts2nestedindex
Outputs a jagged array with same axis-0 shape as counts axis-1
"""
target_offsets = stack.pop()
local_counts = stack.pop()
out = awkward.unflatten(
numpy.arange(target_offsets[-1], dtype=numpy.int64),
awkward.flatten(local_counts),
)
stack.append(out)
def runNN(model, varsIn, varSet, normMean, normStd):
dataset = RootDataset(varsIn=varsIn,
varSet=varSet,
normMean=normMean,
normStd=normStd)
nnOutput = getNNOutput(dataset, model)
fjets = varsIn["fjets"]
counts = ak.num(fjets.pt)
svjJetsAK8 = ak.unflatten(nnOutput, counts)
wpt = 0.5
darksvjJetsAK8 = fjets[svjJetsAK8 >= wpt]
varsIn['nsvjJetsAK8'] = [ak.num(darksvjJetsAK8), "evtw"]
varsIn['nnOutput'] = [svjJetsAK8, "fjw"]
def test_jet_correction_uncertainty():
from coffea.jetmet_tools import JetCorrectionUncertainty
counts, test_eta, test_pt = dummy_jagged_eta_pt()
test_pt_jag = ak.unflatten(test_pt, counts)
test_eta_jag = ak.unflatten(test_eta, counts)
junc_names = ['Summer16_23Sep2016V3_MC_Uncertainty_AK4PFPuppi']
junc = JetCorrectionUncertainty(
**{name: evaluator[name]
for name in junc_names})
print(junc)
juncs = junc.getUncertainty(JetEta=test_eta, JetPt=test_pt)
juncs_jag = list(
junc.getUncertainty(JetEta=test_eta_jag, JetPt=test_pt_jag))
for i, (level, corrs) in enumerate(juncs):
assert (corrs.shape[0] == test_eta.shape[0])
assert (ak.all(corrs == ak.flatten(juncs_jag[i][1])))
def constituents(self, min_pt):
np_results = self._results.to_numpy_with_constituents(min_pt)
off = np.insert(np_results[-1], 0, 0)
out = ak.Array(
ak.layout.ListOffsetArray64(
ak.layout.Index64(np_results[0]), ak.layout.NumpyArray(np_results[1])
)
)
outputs_to_inputs = ak.Array(
ak.layout.ListOffsetArray64(ak.layout.Index64(off), out.layout)
)
shape = ak.num(outputs_to_inputs)
total = np.sum(shape)
duplicate = ak.unflatten(np.zeros(total, np.int64), shape)
prepared = self.data[:, np.newaxis][duplicate]
return prepared[outputs_to_inputs][0]
def get_scale_factor(self, jets, passing_cut):
'''Starting from a jet collection and a string pointing to
the flag defining if the jet is b-tagged or not computes the
per-jet weight to be used. Supports only a single WP for the
moment'''
# First of all flatten everything to make it easier to handle
pt = ak.to_numpy(ak.flatten(jets.pt))
eta = ak.to_numpy(ak.flatten(jets.eta))
flav = ak.to_numpy(ak.flatten(jets.hadronFlavour))
pass_wp = ak.to_numpy(ak.flatten(jets[passing_cut]))
# Get the MC efficiency
eff = self.efficiency_(pt, eta, flav)
# for each systematic/central value compute the proper SF
# cache the SF values as sometimes they are repeated, there
# might also be systematic combinations that are never accessed
# but pruning them at the beginning can be hard
# use schema to define combinations, lcb is a tuple with the sf keys
# for light, charm, bottom for each systematic
flavour_sf_cache = {}
scale_factors = {} # our final product
for key, lcb in self.schema_.items():
# populate cache if needed
for i in range(3):
flavour_sf_cache[lcb[i]] = flavour_sf_cache.get(
# for some reason there is an additional dimension, pass_wp has no effect
lcb[i],
self.sf_[lcb[i]](eta, pt, pass_wp))
scale_factors[key] = eff * self.match_flav_(
flavour_sf_cache[lcb[0]], flavour_sf_cache[lcb[1]],
flavour_sf_cache[lcb[2]], flav)
# use SF and eff to compute p(data) and p(MC)
p_data = {
key: np.where(pass_wp, val, 1 - val)
for key, val in scale_factors.items()
}
p_mc = np.where(pass_wp, eff, 1 - eff)
# return the unflattened version of the ratio
return {
key: ak.unflatten(i / p_mc, ak.num(jets.pt))
for key, i in p_data.items()
}
def constituents(self, min_pt):
self._out = []
self._input_flag = 0
for i in range(len(self._clusterable_level)):
np_results = self._results[i].to_numpy_with_constituents(min_pt)
off = np.insert(np_results[-1], 0, 0)
out = ak.Array(
ak.layout.ListOffsetArray64(
ak.layout.Index64(np_results[0]),
ak.layout.NumpyArray(np_results[1]),
),
behavior=self.data.behavior,
)
outputs_to_inputs = ak.Array(
ak.layout.ListOffsetArray64(ak.layout.Index64(off), out.layout)
)
shape = ak.num(outputs_to_inputs)
total = np.sum(shape)
duplicate = ak.unflatten(np.zeros(total, np.int64), shape)
prepared = self._clusterable_level[i][:, np.newaxis][duplicate]
self._out.append(prepared[outputs_to_inputs])
res = ak.Array(self._replace_multi())
return res
