}
# Set path in docker
ntup_dir = "/data/zprime/ntuples/{}".format(release_name)
arrays_dir = "/data/zprime/arrays/{}".format(release_name)
if not os.path.exists(arrays_dir):
os.makedirs(arrays_dir)
for camp in ["mc16a", "mc16d", "mc16e"]:
# Dump bkg
for index, bkg_key in enumerate(bkg_names):
root_path = ntup_dir + "/" + camp + "/tree_{}.root".format(bkg_names[bkg_key])
dump_flat_ntuple_individual(
root_path,
ntuple_name,
feature_list,
arrays_dir + "/" + camp,
"bkg_{}".format(bkg_key),
use_lower_var_name=True,
)
mz1 = np.load(arrays_dir + "/" + camp + "/bkg_{}_mz1.npy".format(bkg_key))
mz2 = np.load(arrays_dir + "/" + camp + "/bkg_{}_mz2.npy".format(bkg_key))
mz1_mz2 = mz1 - mz2
save_array(mz1_mz2, arrays_dir + "/" + camp + "/", "bkg_{}_mz1_mz2".format(bkg_key))
save_array(
mz1, arrays_dir + "/" + camp + "/", "bkg_{}_mz1_p".format(bkg_key)
) # parameterized feature
save_array(
mz2, arrays_dir + "/" + camp + "/", "bkg_{}_mz2_p".format(bkg_key)
) # parameterized feature
dummy_channel = np.ones(len(mz1))
save_array(dummy_channel, arrays_dir + "/" + camp, "bkg_{}_dummy_channel".format(bkg_key))
arrays_dir = f"/data/zprime/arrays/{release_name}"
if not os.path.exists(arrays_dir):
os.makedirs(arrays_dir)
# Dump bkg
for camp in ["mc16d"]:
# for camp in ["run2"]:
# for camp in ["mc16a", "mc16d", "mc16e"]:
for ntuple_name in bkg_ntuple_names:
for index, bkg_key in enumerate(bkg_names):
root_path = f"{ntup_dir}/{camp}/tree_{bkg_names[bkg_key]}.root"
save_dir_sub = f"{arrays_dir}/{ntuple_name}/{camp}"
dump_flat_ntuple_individual(
root_path,
ntuple_name,
feature_list,
save_dir_sub,
f"bkg_{bkg_key}",
use_lower_var_name=True,
)
mz1 = np.load(f"{save_dir_sub}/bkg_{bkg_key}_mz1.npy")
mz2 = np.load(f"{save_dir_sub}/bkg_{bkg_key}_mz2.npy")
mz1_mz2 = mz1 - mz2
save_array(mz1_mz2, save_dir_sub, f"bkg_{bkg_key}_mz1_mz2")
# parameterized feature
save_array(mz1, save_dir_sub, f"bkg_{bkg_key}_mz1_p")
# parameterized feature
save_array(mz2, save_dir_sub, f"bkg_{bkg_key}_mz2_p")
dummy_channel = np.ones(len(mz1))
save_array(dummy_channel, save_dir_sub,
f"bkg_{bkg_key}_dummy_channel")
normed_weight = np.load(
}
# for cut_level in ["mva_2jet", "mva_2jet_timing"]:
for cut_level in ["mva_2jet"]:
# dump bkg
# ee channel
ee_dir = arrays_dir / "ee" / cut_level
ee_dir.mkdir(parents=True, exist_ok=True)
for bkg_key, bkg_name in bkg_names.items():
root_path = ntup_dir / cut_level / f"{bkg_name}.root"
dump_flat_ntuple_individual(
root_path,
tree_name,
feature_list,
ee_dir,
bkg_name,
use_lower_var_name=True,
channel_feature="is_ee",
)
m_ll = np.load(ee_dir / f"{bkg_name}_ll_m.npy")
# save m_truth
save_array(m_ll, ee_dir, f"{bkg_name}_m_truth")
# mm channel
mm_dir = arrays_dir / "mm" / cut_level
mm_dir.mkdir(parents=True, exist_ok=True)
for bkg_key, bkg_name in bkg_names.items():
root_path = ntup_dir / cut_level / f"{bkg_name}.root"
dump_flat_ntuple_individual(
root_path,
tree_name,