def test_get_jet_names_params():
jet_class = FormJets.Spectral
jet_paramsA = {'MaxMeanDist': 0.4, 'ExpofPTInput': 2.}
floats = np.random.random((3, 8))
# set distance to 0
floats[:, -1] = 0.
for i in range(3):
fill_angular(floats[i])
# need to keep the eventwise file around
with TempTestDir("tst") as dir_name:
eventWise = Components.EventWise(os.path.join(dir_name, "tmp.parquet"))
set_JetInputs(eventWise, floats)
eventWise.selected_event = 0
jet_list = []
with warnings.catch_warnings():
warnings.simplefilter('ignore')
eventWise.selected_event = 0
jets = jet_class(eventWise,
dict_jet_params=jet_paramsA,
run=True,
jet_name="AAJet")
FormJets.check_params(jets, eventWise)
jet_list = [jets]
content = FormJets.create_jet_contents(jet_list, {})
eventWise.append(**content)
jet_paramsB = {k: v for k, v in jet_paramsA.items()}
jet_paramsB['ExpofPTFormatEmbedding'] = None
eventWise.selected_event = 0
jets = jet_class(eventWise,
dict_jet_params=jet_paramsB,
run=True,
jet_name="BBJet")
jet_list = [jets]
content = FormJets.create_jet_contents(jet_list, {})
eventWise.append(**content)
FormJets.check_params(jets, eventWise)
jet_names = FormJets.get_jet_names(eventWise)
assert "AAJet" in jet_names
assert "BBJet" in jet_names
assert len(jet_names) == 2
found_paramsA = FormJets.get_jet_params(eventWise,
"AAJet",
add_defaults=True)
for name in jet_paramsA:
assert found_paramsA[name] == jet_paramsA[name]
found_paramsB = FormJets.get_jet_params(eventWise,
"BBJet",
add_defaults=False)
for name in jet_paramsB:
assert found_paramsB[name] == jet_paramsB[name]
def test_create_jet_contents():
"""Process the unsplit jets in jet_list to from a content dict for an
eventWise.
Parameters
----------
jet_list : list of Agglomerative
the unsplit agglomerative jets, one for each event.
existing_contents : dict
The contents dict for all previous events
Return
------
contents : dict
A dict with the contents of previous and new events.
"""
ints = np.array([[0, 1, -1, -1, 0], [2, 1, -1, -1, 0], [1, -1, 2, 0, 1]])
floats = np.zeros((3, len(FormJets.Agglomerative.float_columns)))
floats[:, 0] = np.arange(3)
agg1 = FormJets.GeneralisedKT((ints, floats),
jet_name="DogJet",
memory_cap=10)
ints = np.array([[0, 1, -1, -1, 0], [2, 1, -1, -1, 0], [1, -1, 2, 0, 1],
[3, -1, -1, -1, 0]])
floats = np.ones((4, len(FormJets.Agglomerative.float_columns)))
agg2 = FormJets.GeneralisedKT((ints, floats),
jet_name="DogJet",
memory_cap=10)
ints = np.empty((0, len(FormJets.Agglomerative.int_columns)))
floats = np.empty((0, len(FormJets.Agglomerative.float_columns)))
agg3 = FormJets.GeneralisedKT((ints, floats),
jet_name="DogJet",
memory_cap=10)
jet_list = [agg1, agg2, agg3]
contents = FormJets.create_jet_contents(jet_list, {})
assert len(contents["DogJet_Label"]) == 3
assert len(contents["DogJet_Label"][0]) == 1
assert len(contents["DogJet_Label"][1]) == 2
assert len(contents["DogJet_Label"][2]) == 0
tst.assert_allclose(sorted(contents["DogJet_PT"][0][0]), np.arange(3))
def solution_to_jet(eventWise, idealised_jet_clusters=None, jet_name=None):
eventWise.selected_event = None
if idealised_jet_clusters is None:
idealised_jet_clusters = eventWise.Solution
if jet_name is None:
jet_name = "SolJet"
# updated_dict will be replaced in the first batch
updated_dict = None
n_events = len(idealised_jet_clusters)
jet_list = []
for event_n, jets_gids in enumerate(idealised_jet_clusters):
if event_n % 100 == 0:
print(f"{event_n/n_events:.1%}", end='\r', flush=True)
eventWise.selected_event = event_n
jet = FormJets.Partitional(eventWise)
source_idx = eventWise.JetInputs_SourceIdx.tolist()
for gids in jets_gids:
if len(gids):
labels = [source_idx.index(i) for i in gids]
jet.create_jet(labels)
jet_list.append(jet)
updated_dict = FormJets.create_jet_contents(jet_list)
eventWise.append(**updated_dict)