2020from ROOT import TF1 , TCanvas , TFile , gStyle
2121
2222from machine_learning_hep .analysis .analyzer import Analyzer
23- from machine_learning_hep .fitting .roofitter import RooFitter
23+ from machine_learning_hep .fitting .roofitter import RooFitter , calc_signif
24+ from machine_learning_hep .fitting .roofitter import create_text_info , add_text_info_fit , add_text_info_perf
2425from machine_learning_hep .utilities import folding , make_message_notfound
2526from machine_learning_hep .utils .hist import (bin_array , create_hist , norm_response , fold_hist ,
2627 fill_hist_fast , get_axis , get_dim , get_bin_limits ,
@@ -71,6 +72,8 @@ def __init__(self, datap, case, typean, period):
7172 self .n_fileresp = os .path .join (self .d_resultsallpmc_proc , self .n_fileresp )
7273 file_result_name = datap ["files_names" ]["resultfilename" ]
7374 self .n_fileresult = os .path .join (self .d_resultsallpdata , file_result_name )
75+ self .p_pdfnames = datap ["analysis" ][self .typean ]['pdf_names' ]
76+ self .p_param_names = datap ["analysis" ][self .typean ]['param_names' ]
7477
7578 self .observables = {
7679 'qa' : ['zg' , 'rg' , 'nsd' , 'zpar' , 'dr' , 'lntheta' , 'lnkt' , 'lntheta-lnkt' ],
@@ -228,7 +231,7 @@ def calculate_efficiencies(self):
228231 self ._save_canvas (c , f'eff/h_ptjet-pthf_eff_{ cat } _ptjet.png' )
229232
230233 # Run 3 efficiencies
231- for cat in cats :
234+ for icat , cat in enumerate ( cats ) :
232235 # gen-level efficiency for feeddown estimation
233236 h_eff_gen = h_genmatch [cat ].Clone ()
234237 h_eff_gen .Divide (h_gen [cat ])
@@ -267,12 +270,23 @@ def calculate_efficiencies(self):
267270 eff = h_det [cat ].Clone (f'h_effnew_{ cat } ' )
268271 ensure_sumw2 (eff )
269272 eff .Divide (h_out )
273+
274+ if eff_corr := self .cfg ('efficiency.reweight' ):
275+ for iptjet in range (get_nbins (eff , 0 )):
276+ for ipt in range (get_nbins (eff , 1 )):
277+ scale_bin (eff , eff_corr [ipt ][icat ], iptjet + 1 , ipt + 1 )
278+
270279 self ._save_hist (eff , f'eff/h_ptjet-pthf_effnew_{ cat } .png' )
271280 self .h_effnew_ptjet_pthf [cat ] = eff
272281
273282 eff_avg = project_hist (h_det [cat ], [1 ], {0 : bins_ptjet })
274283 ensure_sumw2 (eff_avg )
275284 eff_avg .Divide (project_hist (h_out , [1 ], {0 : bins_ptjet }))
285+
286+ if eff_corr := self .cfg ('efficiency.reweight' ):
287+ for ipt in range (get_nbins (eff_avg , 0 )):
288+ scale_bin (eff_avg , eff_corr [ipt ][icat ], ipt + 1 )
289+
276290 self ._save_hist (eff_avg , f'eff/h_pthf_effnew_{ cat } .png' )
277291 self .h_effnew_pthf [cat ] = eff_avg
278292
@@ -341,17 +355,41 @@ def _correct_efficiency(self, hist, ipt):
341355
342356
343357 #region fitting
344- def _roofit_mass (self , hist , ipt , fitcfg , roows = None , filename = None ):
358+ def _roofit_mass (self , level , hist , ipt , pdfnames , param_names , fitcfg , roows = None , filename = None ):
345359 if fitcfg is None :
346360 return None , None
347- res , ws , frame = self .fitter .fit_mass_new (hist , fitcfg , roows , True )
361+ res , ws , frame , residual_frame = self .fitter .fit_mass_new (hist , pdfnames , fitcfg , level , roows , True )
348362 frame .SetTitle (f'inv. mass for p_{{T}} { self .bins_candpt [ipt ]} - { self .bins_candpt [ipt + 1 ]} GeV/c' )
349363 c = TCanvas ()
364+
365+ textInfoRight = create_text_info (0.62 , 0.68 , 1.0 , 0.89 )
366+ add_text_info_fit (textInfoRight , frame , ws , param_names )
367+
368+ textInfoLeft = create_text_info (0.12 , 0.68 , 0.6 , 0.89 )
369+ if level == "data" :
370+ mean_sgn = ws .var (self .p_param_names ["gauss_mean" ])
371+ sigma_sgn = ws .var (self .p_param_names ["gauss_sigma" ])
372+ (sig , sig_err , bkg , bkg_err ,
373+ signif , signif_err , s_over_b , s_over_b_err
374+ ) = calc_signif (ws , res , pdfnames , param_names , mean_sgn , sigma_sgn )
375+
376+ add_text_info_perf (textInfoLeft , sig , sig_err , bkg , bkg_err , s_over_b , s_over_b_err , signif , signif_err )
377+
350378 frame .Draw ()
379+ textInfoRight .Draw ()
380+ textInfoLeft .Draw ()
351381 if res .status () != 0 :
352382 self .logger .warning ('Invalid fit result for %s' , hist .GetName ())
353383 filename = filename .replace ('.png' , '_invalid.png' )
354384 self ._save_canvas (c , filename )
385+
386+ if level == "data" :
387+ residual_frame .SetTitle (f'inv. mass for p_{{T}} { self .bins_candpt [ipt ]} - { self .bins_candpt [ipt + 1 ]} GeV/c' )
388+ cres = TCanvas ()
389+ residual_frame .Draw ()
390+ filename = filename .replace ('.png' , '_residual.png' )
391+ self ._save_canvas (cres , filename )
392+
355393 return res , ws
356394
357395
@@ -493,7 +531,7 @@ def fit(self):
493531 if var := roows .var (par ):
494532 var .setConstant (True )
495533 roo_res , roo_ws = self ._roofit_mass (
496- h_invmass , ipt , fitcfg , roows ,
534+ level , h_invmass , ipt , self . p_pdfnames , self . p_param_names , fitcfg , roows ,
497535 f'roofit/h_mass_fitted{ jetptlabel } _{ string_range_pthf (range_pthf )} _{ level } .png' )
498536 if roo_res .status () != 0 :
499537 self .logger .error ('RooFit failed for %s iptjet %s ipt %d' , level , iptjet , ipt )
@@ -514,8 +552,8 @@ def fit(self):
514552 self .roo_ws_ptjet [level ][jptjet ][ipt ] = roo_ws .Clone ()
515553 # TODO: take parameter names from DB
516554 if level in ('data' , 'mc' ):
517- varname_mean = fitcfg .get ('var_mean' , 'mean' )
518- varname_sigma = fitcfg .get ('var_sigma' , 'sigma_g1' )
555+ varname_mean = fitcfg .get ('var_mean' , self . p_param_names [ "gauss_mean" ] )
556+ varname_sigma = fitcfg .get ('var_sigma' , self . p_param_names [ "gauss_sigma" ] )
519557 self .fit_mean [level ][ipt ] = roo_ws .var (varname_mean ).getValV ()
520558 self .fit_sigma [level ][ipt ] = roo_ws .var (varname_sigma ).getValV ()
521559 varname_m = fitcfg .get ('var' , 'm' )
0 commit comments