77 lines
2.7 KiB
Python
77 lines
2.7 KiB
Python
|
import os
|
||
|
|
import dotenv
|
||
|
|
import sys
|
||
|
|
import argparse
|
||
|
|
import pandas as pd
|
||
|
|
import mplhep
|
||
|
|
import zfit
|
||
|
|
|
||
|
|
from prettytable import PrettyTable
|
||
|
|
|
||
|
|
dotenv.load_dotenv('../properties.env')
|
||
|
|
|
||
|
|
sys.path.insert(0, os.getenv('SYS_PATH'))
|
||
|
|
|
||
|
|
from b2kstll.models.angular import B2Kstll
|
||
|
|
from b2kstll.plot import plot_distributions
|
||
|
|
|
||
|
|
from hep_analytics.processing.extract import FileManager
|
||
|
|
from hep_analytics.processing.transform import select_feature
|
||
|
|
|
||
|
|
FILE_MC = os.getenv('MC_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, tree = "Events", branches = ["q2", "costhetak", "costhetal", "phi"])
|
||
|
|
mc_data = filemanager.extract_data()
|
||
|
|
q2_mc, theta_k_mc, theta_l_mc, phi_mc = mc_data[0], mc_data[1], mc_data[2], mc_data[3]
|
||
|
|
q2_mc, indices = select_feature(feature = q2_mc, limits = bin_ranges[Q2BIN])
|
||
|
|
phi_mc = phi_mc[indices]
|
||
|
|
theta_l_mc = theta_l_mc[indices]
|
||
|
|
theta_k_mc = theta_k_mc[indices]
|
||
|
|
lower_costhetak_cut = float(os.getenv('LOWER_COSTHETAK_CUT'))
|
||
|
|
upper_costhetak_cut = float(os.getenv('UPPER_COSTHETAK_CUT'))
|
||
|
|
theta_k_mc, indices = select_feature(feature = theta_k_mc, limits = (lower_costhetak_cut, upper_costhetak_cut))
|
||
|
|
q2_mc = q2_mc[indices]
|
||
|
|
phi_mc = phi_mc[indices]
|
||
|
|
theta_l_mc = theta_l_mc[indices]
|
||
|
|
|
||
|
|
angular_data = pd.DataFrame({'ctl': theta_l_mc, 'ctk': theta_k_mc, 'phi': phi_mc})
|
||
|
|
angular_data.to_csv(f"ang_fit_mc_{Q2BIN}_bin.csv", index = False)
|
||
|
|
x = B2Kstll('ctl','ctk','phi')
|
||
|
|
x.set_acc(f"./acc_3d_JpsiKstMC_{Q2BIN}_bin.yaml")
|
||
|
|
obs, pdf, params = x.get_pdf('PWave')
|
||
|
|
|
||
|
|
datazfit = zfit.Data.from_pandas(df = angular_data, obs = obs)
|
||
|
|
nll = zfit.loss.UnbinnedNLL(model = pdf, data = datazfit)
|
||
|
|
minimizer = zfit.minimize.Minuit()
|
||
|
|
result = minimizer.minimize(nll)
|
||
|
|
|
||
|
|
param_errors, _ = result.errors()
|
||
|
|
print(param_errors)
|
||
|
|
|
||
|
|
info_table = PrettyTable(["Variable", "Value", "Lower Error", "Upper Error"])
|
||
|
|
fit_labels = ["AFB", "FL", "S3", "S4", "S5", "S7", "S8", "S9"]
|
||
|
|
for name in fit_labels:
|
||
|
|
value = result.params[name]["value"]
|
||
|
|
lower = result.params[name]["minuit_minos"]["lower"]
|
||
|
|
upper = result.params[name]["minuit_minos"]["upper"]
|
||
|
|
info_table.add_row([name, value, lower, upper])
|
||
|
|
|
||
|
|
print(info_table)
|
||
|
|
|
||
|
|
plot_distributions(result, suffix = f"accPwavePDF_mc_ang_fit_{Q2BIN}_bin")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|