}

# 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(
Example #3
0
}

# 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,