def test_sequential_stages(cutflow_1, select_2, infile, full_event_range,
tmpdir):
cutflow_2 = stage.CutFlow("cutflow_2",
str(tmpdir),
selection=select_2,
weights="EventWeight")
chunk = FakeBEEvent(
MaskedUprootTree(infile, event_ranger=full_event_range), "data")
cutflow_1.event(chunk)
cutflow_2.event(chunk)
assert len(chunk.tree) == 2
jet_py = chunk.tree.array("Jet_Py")
assert pytest.approx(
jet_py.flatten()) == [49.641838, 45.008915, -78.01798, 60.730812]
def one_stage_many_calls(selection, tmpdir, chunk_data, chunk_mc):
cutflow_a = stage.CutFlow("cutflow_a",
str(tmpdir),
selection=selection,
weights="EventWeight")
cutflow_b = stage.CutFlow("cutflow_b",
str(tmpdir),
selection=selection,
weights="EventWeight")
cutflow_a.event(chunk_mc)
chunk_mc.tree.reset_mask()
cutflow_a.event(chunk_mc)
cutflow_b.event(chunk_data)
chunk_data.tree.reset_mask()
cutflow_b.event(chunk_data)
collector = cutflow_a.collector()
dataset_readers_list = (
("test_mc", (cutflow_a, )),
("test_data", (cutflow_b, )),
)
return collector, dataset_readers_list
def many_stages_one_call(selection, tmpdir, chunk_data, chunk_mc):
cutflows = [
stage.CutFlow("cutflow_%d" % i,
str(tmpdir),
selection=selection,
weights="EventWeight") for i in range(4)
]
cutflows[0].event(chunk_mc)
chunk_mc.tree.reset_mask()
cutflows[1].event(chunk_mc)
cutflows[2].event(chunk_data)
chunk_data.tree.reset_mask()
cutflows[3].event(chunk_data)
collector = cutflows[0].collector()
dataset_readers_list = (
("test_mc", cutflows[:2]),
("test_data", cutflows[2:]),
)
return collector, dataset_readers_list
def cutflow_1(tmpdir):
return stage.CutFlow("cutflow_1",
str(tmpdir),
selection="NMuon > 1",
weights="NElectron")