EWP-BplusToKstMuMu-AngAna/Code/Scripts/Python Scripts/MC Fit/angular_acceptance_modeling.py

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

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]}")

    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]
    data_mc_phsp    = [phi_mc_phsp, theta_l_mc_phsp, theta_k_mc_phsp, q2_mc_phsp]
    
    df     = pd.DataFrame({'ctl': data_mc_phsp[1], 'ctk': data_mc_phsp[2], 'phi': data_mc_phsp[0], 'q2': q2_mc_phsp})
    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'], ranges = ranges,
                            legendre_orders = orders, calculate_cov = False, chunk_size = 2000)

    out_file = eff.write_to_disk(f'acc_3d_JpsiKstMC_{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']:
        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_phsp_mc_{v}_{Q2BIN}_bin.pdf')
        plt.close()

    
if __name__ == "__main__":
    main()