def test_lumimask():
lumimask = LumiMask(
"tests/samples/Cert_294927-306462_13TeV_EOY2017ReReco_Collisions17_JSON.txt"
)
# pickle & unpickle
lumimask_pickle = cloudpickle.loads(cloudpickle.dumps(lumimask))
# check same mask keys
keys = lumimask._masks.keys()
assert keys == lumimask_pickle._masks.keys()
# check same mask values
assert all(np.all(lumimask._masks[k] == lumimask_pickle._masks[k]) for k in keys)
runs = np.array([303825, 123], dtype=np.uint32)
lumis = np.array([115, 123], dtype=np.uint32)
for lm in lumimask, lumimask_pickle:
mask = lm(runs, lumis)
print("mask:", mask)
assert mask[0]
assert not mask[1]
# test underlying py_func
py_mask = np.zeros(dtype="bool", shape=runs.shape)
LumiMask._apply_run_lumi_mask_kernel.py_func(lm._masks, runs, lumis, py_mask)
assert np.all(mask == py_mask)
assert np.all(lumimask(runs, lumis) == lumimask_pickle(runs, lumis))
def test_export1d():
import uproot
import os
from coffea.hist import export1d
counts, test_eta, test_pt = dummy_jagged_eta_pt()
h_regular_bins = hist.Hist("regular_joe", hist.Bin("x", "x", 20, 0, 200))
h_regular_bins.fill(x=test_pt)
hout = export1d(h_regular_bins)
filename = 'test_export1d.root'
with uproot.create(filename) as fout:
fout['regular_joe'] = hout
fout.close()
with uproot.open(filename) as fin:
hin = fin['regular_joe']
assert (np.all(hin.edges == hout.edges))
assert (np.all(hin.values == hout.values))
del hin
del fin
if os.path.exists(filename):
os.remove(filename)
def test_lumidata():
from numba import types
from numba.typed import Dict
lumidata = LumiData("tests/samples/lumi_small.csv")
# pickle & unpickle
lumidata_pickle = cloudpickle.loads(cloudpickle.dumps(lumidata))
# check same internal lumidata
assert np.all(lumidata._lumidata == lumidata_pickle._lumidata)
runslumis = np.zeros((10, 2), dtype=np.uint32)
results = {"lumi": {}, "index": {}}
for ld in lumidata, lumidata_pickle:
runslumis[:, 0] = ld._lumidata[0:10, 0]
runslumis[:, 1] = ld._lumidata[0:10, 1]
lumi = ld.get_lumi(runslumis)
results["lumi"][ld] = lumi
diff = abs(lumi - 1.539941814)
print("lumi:", lumi, "diff:", diff)
assert diff < 1e-4
# test build_lumi_table_kernel
py_index = Dict.empty(
key_type=types.Tuple([types.uint32, types.uint32]), value_type=types.float64
)
pyruns = ld._lumidata[:, 0].astype("u4")
pylumis = ld._lumidata[:, 1].astype("u4")
LumiData._build_lumi_table_kernel.py_func(
pyruns, pylumis, ld._lumidata, py_index
)
assert len(py_index) == len(ld.index)
# test get_lumi_kernel
py_tot_lumi = np.zeros((1,), dtype=np.float64)
LumiData._get_lumi_kernel.py_func(
runslumis[:, 0], runslumis[:, 1], py_index, py_tot_lumi
)
assert abs(py_tot_lumi[0] - lumi) < 1e-4
# store results:
results["lumi"][ld] = lumi
results["index"][ld] = ld.index
assert np.all(results["lumi"][lumidata] == results["lumi"][lumidata_pickle])
assert len(results["index"][lumidata]) == len(results["index"][lumidata_pickle])
def test_weights():
from coffea.processor import Weights
counts, test_eta, test_pt = dummy_jagged_eta_pt()
scale_central = np.random.normal(loc=1.0, scale=0.01, size=counts.size)
scale_up = scale_central * 1.10
scale_down = scale_central * 0.95
scale_up_shift = 0.10 * scale_central
scale_down_shift = 0.05 * scale_central
weight = Weights(counts.size)
weight.add('test', scale_central, weightUp=scale_up, weightDown=scale_down)
weight.add('testShift',
scale_central,
weightUp=scale_up_shift,
weightDown=scale_down_shift,
shift=True)
var_names = weight.variations
expected_names = ['testShiftUp', 'testShiftDown', 'testUp', 'testDown']
for name in expected_names:
assert (name in var_names)
test_central = weight.weight()
exp_weight = scale_central * scale_central
assert (np.all(np.abs(test_central - (exp_weight)) < 1e-6))
test_up = weight.weight('testUp')
exp_up = scale_central * scale_central * 1.10
assert (np.all(np.abs(test_up - (exp_up)) < 1e-6))
test_down = weight.weight('testDown')
exp_down = scale_central * scale_central * 0.95
assert (np.all(np.abs(test_down - (exp_down)) < 1e-6))
test_shift_up = weight.weight('testUp')
assert (np.all(np.abs(test_shift_up - (exp_up)) < 1e-6))
test_shift_down = weight.weight('testDown')
assert (np.all(np.abs(test_shift_down - (exp_down)) < 1e-6))
def test_issue_247():
from coffea import hist
h = hist.Hist('stuff', hist.Bin('old', 'old', 20, -1, 1))
h.fill(old=h.axis('old').centers())
h2 = h.rebin(h.axis('old'), hist.Bin('new', 'new', 10, -1, 1))
# check first if its even possible to have correct binning
assert np.all(h2.axis('new').edges() == h.axis('old').edges()[::2])
# make sure the lookup works properly
assert np.all(h2.values()[()] == 2.)
h3 = h.rebin(h.axis('old'), 2)
assert np.all(h3.values()[()] == 2.)
with pytest.raises(ValueError):
# invalid division
_ = h.rebin(h.axis('old'), hist.Bin('new', 'new', 8, -1, 1))
newaxis = hist.Bin('new', 'new', h.axis('old').edges()[np.cumsum([0, 2, 3, 5])])
h4 = h.rebin('old', newaxis)
def test_packed_selection():
from coffea.processor import PackedSelection
sel = PackedSelection()
counts, test_eta, test_pt = dummy_jagged_eta_pt()
all_true = np.full(shape=counts.shape, fill_value=True, dtype=np.bool)
all_false = np.full(shape=counts.shape, fill_value=False, dtype=np.bool)
ones = np.ones(shape=counts.shape, dtype=np.uint64)
wrong_shape = ones = np.ones(shape=(counts.shape[0] - 5, ), dtype=np.bool)
sel.add('all_true', all_true)
sel.add('all_false', all_false)
assert (np.all(sel.require(all_true=True, all_false=False) == all_true))
assert (np.all(sel.all('all_true', 'all_false') == all_false))
try:
sel.require(all_true=1, all_false=0)
except ValueError:
pass
try:
sel.add('wrong_shape', wrong_shape)
except ValueError:
pass
try:
sel.add('ones', ones)
except ValueError:
pass
try:
overpack = PackedSelection()
for i in range(65):
overpack.add('sel_%d', all_true)
except RuntimeError:
pass
def test_hist_compat():
from coffea.util import load
test = load('tests/samples/old_hist_format.coffea')
expected_bins = np.array([ -np.inf, 0., 20., 40., 60., 80., 100., 120., 140.,
160., 180., 200., 220., 240., 260., 280., 300., 320.,
340., 360., 380., 400., 420., 440., 460., 480., 500.,
520., 540., 560., 580., 600., 620., 640., 660., 680.,
700., 720., 740., 760., 780., 800., 820., 840., 860.,
880., 900., 920., 940., 960., 980., 1000., 1020., 1040.,
1060., 1080., 1100., 1120., 1140., 1160., 1180., 1200., np.inf,
np.nan])
assert np.all(test._axes[2]._interval_bins[:-1] == expected_bins[:-1])
assert np.isnan(test._axes[2]._interval_bins[-1])
def test_weights_partial():
from coffea.processor import Weights
counts, _, _ = dummy_jagged_eta_pt()
w1 = np.random.normal(loc=1.0, scale=0.01, size=counts.size)
w2 = np.random.normal(loc=1.3, scale=0.05, size=counts.size)
weights = Weights(counts.size, storeIndividual=True)
weights.add('w1', w1)
weights.add('w2', w2)
test_exclude_none = weights.weight()
assert (np.all(np.abs(test_exclude_none - w1 * w2) < 1e-6))
test_exclude1 = weights.partial_weight(exclude=['w1'])
assert (np.all(np.abs(test_exclude1 - w2) < 1e-6))
test_include1 = weights.partial_weight(include=['w1'])
assert (np.all(np.abs(test_include1 - w1) < 1e-6))
test_exclude2 = weights.partial_weight(exclude=['w2'])
assert (np.all(np.abs(test_exclude2 - w1) < 1e-6))
test_include2 = weights.partial_weight(include=['w2'])
assert (np.all(np.abs(test_include2 - w2) < 1e-6))
test_include_both = weights.partial_weight(include=['w1', 'w2'])
assert (np.all(np.abs(test_include_both - w1 * w2) < 1e-6))
# Check that exception is thrown if arguments are incompatible
error_raised = False
try:
weights.partial_weight(exclude=['w1'], include=['w2'])
except ValueError:
error_raised = True
assert (error_raised)
error_raised = False
try:
weights.partial_weight()
except ValueError:
error_raised = True
assert (error_raised)
# Check that exception is thrown if individual weights
# are not saved from the start
weights = Weights(counts.size, storeIndividual=False)
weights.add('w1', w1)
weights.add('w2', w2)
error_raised = False
try:
weights.partial_weight(exclude=['test'], include=['test'])
except ValueError:
error_raised = True
assert (error_raised)
def test_lumimask():
lumimask = LumiMask(
"tests/samples/Cert_294927-306462_13TeV_EOY2017ReReco_Collisions17_JSON.txt"
)
runs = np.array([303825, 123], dtype=np.uint32)
lumis = np.array([115, 123], dtype=np.uint32)
mask = lumimask(runs, lumis)
print("mask:", mask)
assert (mask[0] == True)
assert (mask[1] == False)
# test underlying py_func
py_mask = np.zeros(dtype='bool', shape=runs.shape)
LumiMask.apply_run_lumi_mask_kernel.py_func(lumimask._masks, runs, lumis,
py_mask)
assert (np.all(mask == py_mask))
def test_preloaded_dataframe_getattr():
import uproot
from coffea.processor import PreloadedDataFrame
tree = uproot.open(osp.abspath('tests/samples/nano_dy.root'))['Events']
chunksize = 20
index = 0
arrays = tree.arrays()
df = PreloadedDataFrame(arrays[b'nMuon'].size, arrays)
assert (len(arrays) == len(df))
df['nMuon'] = arrays[b'nMuon']
assert ('nMuon' in df.available)
assert (np.all(df.nMuon == arrays[b'nMuon']))
assert (b'Muon_eta' in df.available)
assert ('nMuon' in df.materialized)
assert (df.size == arrays[b'nMuon'].size)
def test_jec_txt_scalefactors():
extractor = lookup_tools.extractor()
extractor.add_weight_sets([
"testJEC * tests/samples/Fall17_17Nov2017_V32_MC_L2Relative_AK4PFPuppi.jec.txt",
"* * tests/samples/Summer16_07Aug2017_V11_L1fix_MC_L2Relative_AK4PFchs.jec.txt.gz",
"* * tests/samples/Fall17_17Nov2017_V32_MC_Uncertainty_AK4PFPuppi.junc.txt",
"* * tests/samples/Autumn18_V8_MC_UncertaintySources_AK4PFchs.junc.txt",
"* * tests/samples/Spring16_25nsV10_MC_SF_AK4PFPuppi.jersf.txt",
"* * tests/samples/Autumn18_V7b_MC_SF_AK8PFchs.jersf.txt.gz",
"* * tests/samples/Fall17_17Nov2017_V32_MC_L2Relative_AK4Calo.jec.txt.gz",
"* * tests/samples/Fall17_17Nov2017_V32_MC_L1JPTOffset_AK4JPT.jec.txt.gz",
"* * tests/samples/Fall17_17Nov2017B_V32_DATA_L2Relative_AK4Calo.txt.gz",
"* * tests/samples/Autumn18_V7b_DATA_SF_AK4PF.jersf.txt",
"* * tests/samples/Autumn18_RunC_V19_DATA_L2Relative_AK8PFchs.jec.txt.gz",
"* * tests/samples/Autumn18_RunA_V19_DATA_L2Relative_AK4Calo.jec.txt",
])
extractor.finalize()
evaluator = extractor.make_evaluator()
counts, test_eta, test_pt = dummy_jagged_eta_pt()
# test structured eval
test_eta_jagged = ak.unflatten(test_eta, counts)
test_pt_jagged = ak.unflatten(test_pt, counts)
jec_out = evaluator[
"testJECFall17_17Nov2017_V32_MC_L2Relative_AK4PFPuppi"](test_eta,
test_pt)
jec_out_jagged = evaluator[
"testJECFall17_17Nov2017_V32_MC_L2Relative_AK4PFPuppi"](
test_eta_jagged, test_pt_jagged)
print(evaluator["testJECFall17_17Nov2017_V32_MC_L2Relative_AK4PFPuppi"])
jec_out = evaluator["Summer16_07Aug2017_V11_L1fix_MC_L2Relative_AK4PFchs"](
test_eta, test_pt)
jec_out_jagged = evaluator[
"Summer16_07Aug2017_V11_L1fix_MC_L2Relative_AK4PFchs"](test_eta_jagged,
test_pt_jagged)
print(jec_out)
print(jec_out_jagged)
print(evaluator["Summer16_07Aug2017_V11_L1fix_MC_L2Relative_AK4PFchs"])
jersf_out = evaluator["Spring16_25nsV10_MC_SF_AK4PFPuppi"](test_eta,
test_pt)
jersf_out_jagged = evaluator["Spring16_25nsV10_MC_SF_AK4PFPuppi"](
test_eta_jagged, test_pt_jagged)
print(jersf_out)
print(jersf_out_jagged)
# single jet jersf lookup test:
single_jersf_out_1d = evaluator["Spring16_25nsV10_MC_SF_AK4PFPuppi"](
np.array([1.4]), np.array([44.0]))
single_jersf_out_0d = evaluator["Spring16_25nsV10_MC_SF_AK4PFPuppi"](
np.array(1.4), np.array(44.0))
truth_out = np.array([[1.084, 1.095, 1.073]], dtype=np.float32)
assert np.all(single_jersf_out_1d == truth_out)
assert np.all(single_jersf_out_0d == truth_out)
print(evaluator["Spring16_25nsV10_MC_SF_AK4PFPuppi"])
junc_out = evaluator["Fall17_17Nov2017_V32_MC_Uncertainty_AK4PFPuppi"](
test_eta, test_pt)
junc_out_jagged = evaluator[
"Fall17_17Nov2017_V32_MC_Uncertainty_AK4PFPuppi"](test_eta_jagged,
test_pt_jagged)
print(junc_out)
print(junc_out_jagged)
print(evaluator["Fall17_17Nov2017_V32_MC_Uncertainty_AK4PFPuppi"])
assert ("Autumn18_V8_MC_UncertaintySources_AK4PFchs_AbsoluteScale"
in evaluator.keys())
junc_out = evaluator[
"Autumn18_V8_MC_UncertaintySources_AK4PFchs_AbsoluteScale"](test_eta,
test_pt)
junc_out_jagged = evaluator[
"Autumn18_V8_MC_UncertaintySources_AK4PFchs_AbsoluteScale"](
test_eta_jagged, test_pt_jagged)
print(junc_out)
print(junc_out_jagged)
print(
evaluator["Autumn18_V8_MC_UncertaintySources_AK4PFchs_AbsoluteScale"])
def test_issue_333():
axis = hist.Bin("channel", "Channel b1", 50, 0, 2000)
temp = np.arange(0, 2000, 40, dtype=np.int16)
assert np.all(axis.index(temp) == np.arange(50) + 1)