def test__infer_on_cut(dummy_data_training):
momenta, cut_mom, near_mom, labels, cut_labs, near_labs, delta_cut, delta_near = dummy_data_training
fks = FKSPartition(momenta=momenta, labels=labels, all_legs=False)
cut_momenta, near_momenta, cut_labels, near_labels = fks.cut_near_split(
delta_cut, delta_near)
n_gluon = 1
NN = Model((n_gluon + 2 - 1) * 4,
cut_momenta,
cut_labels,
all_jets=False,
all_legs=False)
_, _, _, _, _, _, _, _ = NN.process_training_data()
model_cut, x_mean_cut, x_std_cut, y_mean_cut, y_std_cut = train_cut_network(
cut_momenta=cut_momenta,
NJ_cut=cut_labels,
order='LO',
n_gluon=1,
delta_cut=delta_cut,
points=len(cut_momenta) * 2,
model_dir='',
epochs=1)
y_pred_cuts = infer_on_cut(NN=NN,
moms=cut_momenta,
model=model_cut,
x_mean_cut=x_mean_cut,
x_std_cut=x_std_cut,
y_mean_cut=y_mean_cut,
y_std_cut=y_std_cut)
def test__train_near_networks_general_high_precision(
dummy_data_all_legs_training):
momenta, cut_mom, near_mom, labels, cut_labs, near_labs, delta_cut, delta_near = dummy_data_all_legs_training
nlegs = len(momenta[0]) - 2
fks = FKSPartition(momenta=momenta, labels=labels, all_legs=True)
cut_momenta, near_momenta, cut_labels, near_labels = fks.cut_near_split(
delta_cut, delta_near)
pairs, labs_split = fks.weighting()
model_near, x_mean_near, x_std_near, y_mean_near, y_std_near = train_near_networks_general(
input_size=(nlegs + 2) * 4,
pairs=pairs,
near_momenta=near_momenta,
NJ_split=labs_split,
delta_near=delta_near,
model_dir='',
all_jets=True,
all_legs=True,
model_dataset=False,
epochs=1,
high_precision=True)
def recut_data(self, momenta, nj):
fks = FKSPartition(momenta=momenta, labels=nj, all_legs=self.all_legs)
cut_momenta, near_momenta, cut_nj, near_nj = fks.cut_near_split(
delta_cut=self.delta_cut, delta_near=0.02)
momenta = np.concatenate((cut_momenta, near_momenta))
nj = np.concatenate((cut_nj, near_nj))
return momenta, nj
def split_data(self, momenta, nj, return_weights=False):
fks = FKSPartition(momenta=momenta, labels=nj, all_legs=self.all_legs)
cut_momenta, near_momenta, cut_nj, near_nj = fks.cut_near_split(
delta_cut=self.delta_cut, delta_near=self.delta_near)
if return_weights:
pairs, near_nj_split, weights = fks.weighting(
return_weights=return_weights)
return cut_momenta, near_momenta, cut_nj, near_nj, pairs, near_nj_split, weights
else:
pairs, near_nj_split = fks.weighting(return_weights=return_weights)
return cut_momenta, near_momenta, cut_nj, near_nj, pairs, near_nj_split
def test__weighting_all_legs(dummy_data_all_legs):
momenta_all_legs, cut_mom, near_mom_all_legs, labels, cut_label, near_label, delta_cut, delta_near = dummy_data_all_legs
fks = FKSPartition(
momenta = momenta_all_legs,
labels = labels,
all_legs = True
)
cut_momenta, near_momenta, cut_labels, near_labels = fks.cut_near_split(delta_cut, delta_near)
pairs, labs_split = fks.weighting()
assert len(pairs) == 9
assert len(labs_split) == 9
assert np.isclose(np.sum(labs_split),near_label)
def test__cut_near_split_all_legs(dummy_data_all_legs):
momenta_all_legs, cut_mom, near_mom_all_legs, labels, cut_label, near_label, delta_cut, delta_near = dummy_data_all_legs
fks = FKSPartition(
momenta = momenta_all_legs,
labels = labels,
all_legs = True
)
cut_momenta, near_momenta, cut_labels, near_labels = fks.cut_near_split(delta_cut, delta_near)
assert check_list(cut_momenta[0][2],cut_mom[0][2])
assert check_list(near_momenta[0][2], near_mom_all_legs[0][2])
assert cut_labels[0] == cut_label
assert near_labels[0] == near_label
def test__train_cut_network(dummy_data_training):
momenta, cut_mom, near_mom, labels, cut_labs, near_labs, delta_cut, delta_near = dummy_data_training
fks = FKSPartition(momenta=momenta, labels=labels, all_legs=False)
cut_momenta, near_momenta, cut_labels, near_labels = fks.cut_near_split(
delta_cut, delta_near)
model_cut, x_mean_cut, x_std_cut, y_mean_cut, y_std_cut = train_cut_network(
cut_momenta=cut_momenta,
NJ_cut=cut_labels,
order='LO',
n_gluon=1,
delta_cut=delta_cut,
points=len(cut_momenta) * 2,
model_dir='',
epochs=1)
def create_model_all_legs_dataset(dummy_data_all_legs_training):
momenta, cut_mom, near_mom, labels, cut_labs, near_labs, delta_cut, delta_near = dummy_data_all_legs_training
nlegs = len(momenta[0]) - 2
fks = FKSPartition(momenta=momenta, labels=labels, all_legs=True)
cut_momenta, near_momenta, cut_labels, near_labels = fks.cut_near_split(
delta_cut, delta_near)
pairs, labs_split = fks.weighting()
model = Model(input_size=(nlegs + 2) * 4,
momenta=cut_momenta,
labels=cut_labels,
all_jets=False,
all_legs=True,
model_dataset=True)
return model
def test__train_near_networks(dummy_data_training):
momenta, cut_mom, near_mom, labels, cut_labs, near_labs, delta_cut, delta_near = dummy_data_training
fks = FKSPartition(momenta=momenta, labels=labels, all_legs=False)
cut_momenta, near_momenta, cut_labels, near_labels = fks.cut_near_split(
delta_cut, delta_near)
pairs, labs_split = fks.weighting()
model_near, x_mean_near, x_std_near, y_mean_near, y_std_near = train_near_networks(
pairs=pairs,
near_momenta=near_momenta,
NJ_split=labs_split,
order='LO',
n_gluon=1,
delta_near=delta_near,
points=len(near_momenta) * 2,
model_dir='',
epochs=1)
def create_model_high_precision(dummy_data_training):
momenta, cut_mom, near_mom, labels, cut_labs, near_labs, delta_cut, delta_near = dummy_data_training
nlegs = len(momenta[0]) - 2
n_gluon = 1
fks = FKSPartition(momenta=momenta, labels=labels, all_legs=False)
cut_momenta, near_momenta, cut_labels, near_labels = fks.cut_near_split(
delta_cut, delta_near)
pairs, labs_split = fks.weighting()
model = Model(input_size=(n_gluon + 2 - 1) * 4,
momenta=cut_momenta,
labels=cut_labels,
all_jets=False,
all_legs=False,
high_precision=True)
return model
def test__train_cut_network_general_dataset(dummy_data_all_legs_training):
momenta, cut_mom, near_mom, labels, cut_labs, near_labs, delta_cut, delta_near = dummy_data_all_legs_training
nlegs = len(momenta[0]) - 2
fks = FKSPartition(momenta=momenta, labels=labels, all_legs=True)
cut_momenta, near_momenta, cut_labels, near_labels = fks.cut_near_split(
delta_cut, delta_near)
model_cut, x_mean_cut, x_std_cut, y_mean_cut, y_std_cut = train_cut_network_general(
input_size=(nlegs + 2) * 4,
cut_momenta=cut_momenta,
NJ_cut=cut_labels,
delta_cut=delta_cut,
model_dir='',
all_jets=True,
all_legs=True,
model_dataset=True,
epochs=1)
def test__infer_on_near_splits_separate(dummy_data_training):
momenta, cut_mom, near_mom, labels, cut_labs, near_labs, delta_cut, delta_near = dummy_data_training
fks = FKSPartition(momenta=momenta, labels=labels, all_legs=False)
cut_momenta, near_momenta, cut_labels, near_labels = fks.cut_near_split(
delta_cut, delta_near)
pairs, labs_split = fks.weighting()
n_gluon = 1
NN = Model((n_gluon + 2 - 1) * 4,
near_momenta,
labs_split[0],
all_jets=False,
all_legs=False)
_, _, _, _, _, _, _, _ = NN.process_training_data()
model_near, x_mean_near, x_std_near, y_mean_near, y_std_near = train_near_networks(
pairs=pairs,
near_momenta=near_momenta,
NJ_split=labs_split,
order='LO',
n_gluon=1,
delta_near=delta_near,
points=len(near_momenta) * 2,
model_dir='',
epochs=1)
y_preds_nears, y_pred_nears = infer_on_near_splits_separate(
NN=NN,
moms=near_momenta,
models=model_near,
x_mean_near=x_mean_near,
x_std_near=x_std_near,
y_mean_near=y_mean_near,
y_std_near=y_std_near)