Spaces:
Runtime error
Runtime error
| import os | |
| from tqdm import tqdm | |
| import argparse | |
| import numpy as np | |
| import pandas as pd | |
| import pickle | |
| from src.dataset.get_dataset import get_iter | |
| from src.utils.paths import get_path | |
| from pathlib import Path | |
| import torch | |
| # This script attempts to open dataset files and prints the number of events in each one. | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--input", type=str, required=True) | |
| parser.add_argument("--dataset-cap", type=int, default=-1) | |
| parser.add_argument("--output", type=str, default="") | |
| parser.add_argument("--plot-only", action="store_true") | |
| # Plots of stats: total visible energy, visible mass, number of AK8 jets, number of pfcands + special_pfcands | |
| args = parser.parse_args() | |
| path = get_path(args.input, "preprocessed_data") | |
| if args.output == "": | |
| args.output = args.input | |
| output_path = os.path.join(get_path(args.output, "results"), "dataset_stats") | |
| Path(output_path).mkdir(parents=True, exist_ok=True) | |
| if not args.plot_only: | |
| stats = {} | |
| for subdataset in os.listdir(path): | |
| print("-----", subdataset, "-----") | |
| current_path = os.path.join(path, subdataset) | |
| dataset = get_iter(current_path) | |
| n = 0 | |
| stats[subdataset] = {"total_visible_E": [], "visible_mass": [], "n_fatjets": [], "n_pfcands": [], "pt": [], | |
| "fatjet_pt": [], "genjet_pt": []} | |
| for data in tqdm(dataset): | |
| n += 1 | |
| if args.dataset_cap != -1 and n > args.dataset_cap: | |
| break | |
| n_fatjets = len(data.fatjets) | |
| n_pfcands = len(data.pfcands) #+ len(data.special_pfcands) | |
| total_visible_E = torch.sum(data.pfcands.E) #+ torch.sum(data.special_pfcands.E) | |
| pt = data.pfcands.pt.tolist() | |
| visible_mass = torch.sqrt(torch.sum(data.pfcands.E)**2 - torch.sum(data.pfcands.p)**2) | |
| stats[subdataset]["total_visible_E"].append(total_visible_E) | |
| stats[subdataset]["visible_mass"].append(visible_mass) | |
| stats[subdataset]["n_fatjets"].append(n_fatjets) | |
| stats[subdataset]["n_pfcands"].append(n_pfcands) | |
| stats[subdataset]["pt"] += pt | |
| stats[subdataset]["fatjet_pt"] += data.fatjets.pt.tolist() | |
| stats[subdataset]["genjet_pt"] += data.genjets.pt.tolist() | |
| #stats[subdataset]["n_events"] = dataset.n_events | |
| def get_properties(name): | |
| # get mediator mass, dark quark mass, r_inv from the filename | |
| parts = name.split("_") | |
| mMed = int(parts[1].split("-")[1]) | |
| mDark = int(parts[2].split("-")[1]) | |
| rinv = float(parts[3].split("-")[1]) | |
| return mMed, mDark, rinv | |
| result = {} | |
| for key in stats: | |
| mMed, mDark, rinv = get_properties(key) | |
| if mMed not in result: | |
| result[mMed] = {} | |
| if mDark not in result[mMed]: | |
| result[mMed][mDark] = {} | |
| result[mMed][mDark][rinv] = stats[key] | |
| pickle.dump(result, open(os.path.join(output_path, "result.pkl"), "wb")) | |
| if args.plot_only: | |
| result = pickle.load(open(os.path.join(output_path, "result.pkl"), "rb")) | |
| import matplotlib.pyplot as plt | |
| # heatmap plots | |
| mediator_masses = sorted(list(result.keys())) | |
| dark_masses = [20] | |
| r_invs = sorted(list(set([rinv for mMed in result for mDark in result[mMed] for rinv in result[mMed][mDark]]))) | |
| def plot_distribution(result, key_name): | |
| print("---> key:", key_name) | |
| fig, ax = plt.subplots(len(mediator_masses), len(r_invs), figsize=(3*len(r_invs), 3*len(mediator_masses))) | |
| for i, mMed in enumerate(mediator_masses): | |
| for j, rinv in enumerate(r_invs): | |
| mDark = dark_masses[0] | |
| data = result[mMed][mDark][rinv][key_name] | |
| if key_name == "n_pfcands": | |
| number_of_zeros = len(data) - np.count_nonzero(data) | |
| print(f"Number of zeros in {mMed} {rinv}: {number_of_zeros}") | |
| if key_name == "fatjet_pt": | |
| print("Min fatjet_pt:", min(data)) | |
| if key_name == "genjet_pt": | |
| print("Min genjet_pt:", min(data)) | |
| number_of_zeros = len(data) - np.count_nonzero(data) | |
| print(f"Number of zeros in {mMed} {rinv}: {number_of_zeros}") | |
| ax[i, j].hist(data, bins=50) | |
| ax[i, j].set_title(f"$m_{{Z'}}$={mMed},$r_{{inv}}$={rinv} ($\Sigma$={int(sum(data))})") | |
| if key_name == "pt": | |
| ax[i, j].set_yscale("log") | |
| # big title | |
| fig.suptitle(key_name) | |
| fig.tight_layout() | |
| fig.savefig(os.path.join(output_path, f"{key_name}.pdf")) | |
| #fig.show() | |
| plot_distribution(result, "total_visible_E") | |
| plot_distribution(result, "visible_mass") | |
| plot_distribution(result, "n_fatjets") | |
| plot_distribution(result, "n_pfcands") | |
| plot_distribution(result, "pt") | |
| plot_distribution(result, "fatjet_pt") | |
| plot_distribution(result, "genjet_pt") | |