95 lines
4.5 KiB
Python
95 lines
4.5 KiB
Python
|
import os
|
||
|
|
import dotenv
|
||
|
|
import sys
|
||
|
|
import argparse
|
||
|
|
import mplhep
|
||
|
|
import numpy as np
|
||
|
|
import pandas as pd
|
||
|
|
import matplotlib.pyplot as plt
|
||
|
|
|
||
|
|
dotenv.load_dotenv('../properties.env')
|
||
|
|
|
||
|
|
sys.path.insert(0, os.getenv('SYS_PATH'))
|
||
|
|
|
||
|
|
from analysis.efficiency import get_efficiency_model_class
|
||
|
|
|
||
|
|
from hep_analytics.processing.extract import FileManager
|
||
|
|
from hep_analytics.processing.transform import select_feature, reweight_feature
|
||
|
|
from hep_analytics.processing.visualisation import reweight_comparing_plot
|
||
|
|
|
||
|
|
FILE_GEN = os.getenv('GEN_FILE')
|
||
|
|
FILE_MC_PHSP = os.getenv('MC_PHSP_FILE')
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
parser = argparse.ArgumentParser()
|
||
|
|
parser.add_argument('--q2bin', dest ='q2bin', default = 0)
|
||
|
|
args = parser.parse_args()
|
||
|
|
Q2BIN = int(args.q2bin)
|
||
|
|
|
||
|
|
mplhep.style.use("LHCb2")
|
||
|
|
|
||
|
|
bin_ranges = [(0.25, 4.00), (4.00, 8.00), (11.00, 12.50), (15.00, 18.00), (1.10, 6.00), (1.1, 2.5), (2.5, 4.0), (4.0, 6.0), (6.0, 8.0)]
|
||
|
|
print(f"Selected Q2 Bin Range is {bin_ranges[Q2BIN]}")
|
||
|
|
|
||
|
|
bin_labels = [r"$0.25 < q^2 < 4.00$", r"$4.00 < q^2 < 8.00$", r"$11.00 < q^2 < 12.50$",
|
||
|
|
r"$15.00 < q^2 < 18.00$", r"$1.10 < q^2 < 6.00$",
|
||
|
|
r"$1.1 < q^2 < 2.5$", r"$2.5 < q^2 < 4.0$", r"$4.0 < q^2 < 6.0$", r"$6.0 < q^2 < 8.0$"]
|
||
|
|
|
||
|
|
filemanager = FileManager(file = FILE_MC_PHSP, tree = "Events", branches = ["q2", "costhetak", "costhetal", "phi"])
|
||
|
|
mc_phsp_data = filemanager.extract_data()
|
||
|
|
q2_mc_phsp, theta_k_mc_phsp, theta_l_mc_phsp, phi_mc_phsp = mc_phsp_data[0], mc_phsp_data[1], mc_phsp_data[2], mc_phsp_data[3]
|
||
|
|
q2_mc_phsp, indices = select_feature(feature = q2_mc_phsp, limits = bin_ranges[Q2BIN])
|
||
|
|
phi_mc_phsp = phi_mc_phsp[indices]
|
||
|
|
theta_l_mc_phsp = theta_l_mc_phsp[indices]
|
||
|
|
theta_k_mc_phsp = theta_k_mc_phsp[indices]
|
||
|
|
lower_costhetak_cut = float(os.getenv('LOWER_COSTHETAK_CUT'))
|
||
|
|
upper_costhetak_cut = float(os.getenv('UPPER_COSTHETAK_CUT'))
|
||
|
|
theta_k_mc_phsp, indices = select_feature(feature = theta_k_mc_phsp, limits = (lower_costhetak_cut, upper_costhetak_cut))
|
||
|
|
q2_mc_phsp = q2_mc_phsp[indices]
|
||
|
|
phi_mc_phsp = phi_mc_phsp[indices]
|
||
|
|
theta_l_mc_phsp = theta_l_mc_phsp[indices]
|
||
|
|
|
||
|
|
filemanager = FileManager(file = FILE_GEN, tree = "Events", branches = ["q2", "costhetak", "costhetal", "phi"])
|
||
|
|
gen_data = filemanager.extract_data()
|
||
|
|
q2_gen, theta_k_gen, theta_l_gen, phi_gen = gen_data[0], gen_data[1], gen_data[2], gen_data[3]
|
||
|
|
q2_gen, indices = select_feature(feature = q2_gen, limits = bin_ranges[Q2BIN])
|
||
|
|
phi_gen = phi_gen[indices]
|
||
|
|
theta_l_gen = theta_l_gen[indices]
|
||
|
|
theta_k_gen = theta_k_gen[indices]
|
||
|
|
theta_k_gen, indices = select_feature(feature = theta_k_gen, limits = (lower_costhetak_cut, upper_costhetak_cut))
|
||
|
|
q2_gen = q2_gen[indices]
|
||
|
|
phi_gen = phi_gen[indices]
|
||
|
|
theta_l_gen = theta_l_gen[indices]
|
||
|
|
|
||
|
|
q2_mc_phsp_weights = reweight_feature(original_feature = q2_mc_phsp, target_feature = q2_gen, n_bins = 25)
|
||
|
|
reweight_comparing_plot(original_feature = q2_mc_phsp, target_feature = q2_gen, weights = q2_mc_phsp_weights,
|
||
|
|
n_bins = 25, suptitle = f"{bin_labels[Q2BIN]}",
|
||
|
|
titles = ["Q2 MC PHSP", "Q2 MC PHSP (reweighted)", "Q2 Gen"], save = f"reweighted_q2_gen_mc_phsp_{Q2BIN}_bin.png")
|
||
|
|
|
||
|
|
df = pd.DataFrame({'ctl': theta_l_mc_phsp, 'ctk': theta_k_mc_phsp, 'phi': phi_mc_phsp, 'q2': q2_mc_phsp, 'weights': q2_mc_phsp_weights})
|
||
|
|
orders = {"ctl": 4, "ctk": 6, "phi": 2}
|
||
|
|
ranges = {"ctl": [-1.0, 1.0], "ctk": [lower_costhetak_cut, upper_costhetak_cut], "phi": [-np.pi, np.pi]}
|
||
|
|
EffClass = get_efficiency_model_class('legendre')
|
||
|
|
eff = EffClass.fit(df, ['ctl', 'ctk', 'phi'], weight_var = "weights", ranges = ranges,
|
||
|
|
legendre_orders = orders, calculate_cov = False, chunk_size = 2000)
|
||
|
|
|
||
|
|
out_file = eff.write_to_disk(f'acc_3d_JpsiKstMC_reweighted_{Q2BIN}_bin.yaml')
|
||
|
|
print(out_file)
|
||
|
|
|
||
|
|
labels = {'ctl': r'$\cos \theta_L$', 'ctk': r'$\cos \theta_K$', 'phi': '$\phi$', 'q2': '$q^2$ [GeV$^2$]'}
|
||
|
|
for v in ['ctl', 'ctk', 'phi']:
|
||
|
|
fig, ax = plt.subplots(figsize = (15, 10))
|
||
|
|
plt.xlim(*ranges[v])
|
||
|
|
x, y = eff.project_efficiency(v, n_points = 1000)
|
||
|
|
plt.plot(x, y, 'b-')
|
||
|
|
plt.hist(df[v], density = True, bins = 50, color = 'grey', alpha = 0.5)
|
||
|
|
plt.ylabel("a.u.", horizontalalignment = 'right', y = 1.0)
|
||
|
|
plt.xlabel(labels[v], horizontalalignment = 'right', x = 1.0)
|
||
|
|
plt.savefig(f'acc_3d_JpsiKstMC_{v}_{Q2BIN}_bin.pdf')
|
||
|
|
plt.close()
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|