def cluster_sequence(file_name,
names,
classes,
params,
jet_pt_cut,
dijet_mass,
pileup_file=None,
mean_pileup_per_event=50,
z_cuts=[0.2, 0.4]):
if pileup_file is not None:
ew = NonwritingEventWise.from_file(file_name)
else:
ew = Components.EventWise.from_file(file_name)
if pileup_file is not None:
print("Adding pileup")
ew = remove_excess_data(ew)
pileup = NonwritingEventWise.from_file(pileup_file)
ew = AddPileup.add_pileup(ew, pileup, mean_pileup_per_event)
filter_functions = "remove pileup"
else:
filter_functions = "ignore pileup"
print("Adding jet inputs")
FormJetInputs.create_jetInputs(ew, filter_functions=filter_functions)
print("Adding jets")
for name, cla, param in zip(names, classes, params):
finished = False
while not finished:
finished = FormJets.cluster_multiapply(ew, cla, param, name,
np.inf)
n_jets = len(names) * 20
if z_cuts:
print("Doing z_cuts")
ParallelFormJets.batch_masks(ew, z_cuts, n_jets)
ParallelFormJets.batch_filter(ew, n_jets)
print("Calculating scores")
if 'top' in str(dijet_mass):
ParallelFormJets.batch_semileptonic(ew, n_jets)
ParallelFormJets.batch_correct_semileptonic_masses(ew, n_jets)
else:
CompareClusters.append_scores(ew,
dijet_mass=dijet_mass,
end_time=np.inf)
ParallelFormJets.batch_correct_masses(ew, jet_pt_cut, n_jets)
return ew
def test_append_scores():
# set up as for filter jets
params = {}
# event 0
params['Jet_Label'] = [ak.from_iter([])]
params['Jet_Parent'] = [ak.from_iter([])]
params['Jet_Child1'] = [ak.from_iter([])]
params['Jet_PT'] = [ak.from_iter([])]
params['Energy'] = [ak.from_iter([])]
params['Px'] = [ak.from_iter([])]
params['Py'] = [ak.from_iter([])]
params['Pz'] = [ak.from_iter([])]
params['Children'] = [ak.from_iter([])]
params['Parents'] = [ak.from_iter([])]
params['MCPID'] = [ak.from_iter([])]
params['JetInputs_SourceIdx'] = [ak.from_iter([])]
params['TagIndex'] = [ak.from_iter([])]
# event 1
params['Jet_Label'] = [ak.from_iter([[0], [1, 2, 3, 4, 5]])]
params['Jet_Parent'] += [ak.from_iter([[-1], [-1, 1, 1, 2, 2]])]
params['Jet_Child1'] += [ak.from_iter([[-1], [1, 2, -1, -1, -1]])]
params['Jet_PT'] += [ak.from_iter([[
50.,
], [0.2, 0.1, 0., 0., .1]])]
params['JetInputs_SourceIdx'] += [ak.from_iter(np.arange(6))]
params['Energy'] += [
ak.from_iter([30., 10., 20., 70., 20., 10., 45., 56., 40., 25.])
]
params['Px'] += [ak.from_iter([3., 0., 2., 1., 2., -1., 0., 3., -1., 0.])]
params['Py'] += [ak.from_iter([3., 0., 2., 2., 2., 1., -1., -3., 0., -1.])]
params['Pz'] += [ak.from_iter([3., 0., 2., 0., 2., 2., -5., -2., 1., 0.])]
params['Children'] += [
ak.from_iter([[], [3], [], [5], [], [], [2, 7, 8, 9], [], [], [], []])
]
params['Parents'] += [
ak.from_iter([[], [], [6], [1], [], [3], [], [6], [6], [6], []])
]
params['MCPID'] += [ak.from_iter([4, -5, 5, 3, 2, 1, -5, -1, 7, 11, 12])]
params['TagIndex'] += [ak.from_iter([1, 6])]
# event 2
params['Jet_Label'] = [ak.from_iter([[0], [1, 2, 3]])]
params['Jet_Parent'] += [ak.from_iter([[-1], [-1, 1, 1]])]
params['Jet_Child1'] += [ak.from_iter([[-1], [1, -1, -1]])]
params['Jet_PT'] += [ak.from_iter([[
50.,
], [21., 0.1, 0.]])]
params['JetInputs_SourceIdx'] += [ak.from_iter(np.arange(6))]
params['Energy'] += [
ak.from_iter([30., 10., 20., 70., 20., 10., 45., 56., 40., 25.])
]
params['Px'] += [ak.from_iter([3., 0., 2., 1., 2., -1., 0., 3., -1., 0.])]
params['Py'] += [ak.from_iter([3., 0., 2., 2., 2., 1., -1., -3., 0., -1.])]
params['Pz'] += [ak.from_iter([3., 0., 2., 0., 2., 2., -5., -2., 1., 0.])]
params['Children'] += [
ak.from_iter([[], [3], [], [5], [], [], [2, 7, 8, 9], [], [], [], []])
]
params['Parents'] += [
ak.from_iter([[], [], [6], [1], [], [3], [], [6], [6], [6], []])
]
params['MCPID'] += [ak.from_iter([4, -5, 5, 3, 2, 1, -5, -1, 7, 11, 12])]
params['TagIndex'] += [ak.from_iter([1, 6])]
params = {key: ak.from_iter(val) for key, val in params.items()}
with TempTestDir("tst") as dir_name:
# this will raise a value error if given an empty eventWise
file_name = "test.parquet"
ew = Components.EventWise(os.path.join(dir_name, file_name))
ew.append(**params)
# mock JetQuality.quality_width_fracton and get_detectable_comparisons
with unittest.mock.patch('jet_tools.JetQuality.quality_width_fracton',
new=fake_quality_width_fraction):
with unittest.mock.patch(
'jet_tools.CompareClusters.get_detectable_comparisons',
new=fake_detectable_comparisons):
with unittest.mock.patch(
'jet_tools.TrueTag.add_detectable_fourvector',
new=fake_empty):
import jet_tools
CompareClusters.append_scores(ew)
tst.assert_allclose(ew.Jet_Bork, [np.inf, np.inf, np.inf])
tst.assert_allclose(ew.Jet_Quack[0], [0.5, np.nan])
tst.assert_allclose(ew.Jet_Quack[1], [0.5, -np.inf])
tst.assert_allclose(ew.Jet_Quack[2], [])
assert np.isnan(ew.Jet_AveBork)
tst.assert_allclose(ew.Jet_AveQuack, 0.5)
assert np.isnan(ew.Jet_QualityWidth)
tst.assert_allclose(ew.Jet_QualityFraction,
3 / Constants.dijet_mass)