diff --git a/PWGHF/D2H/Macros/hf_analysis_utils.py b/PWGHF/D2H/Macros/hf_analysis_utils.py index bc0528f38ef..862dce80717 100644 --- a/PWGHF/D2H/Macros/hf_analysis_utils.py +++ b/PWGHF/D2H/Macros/hf_analysis_utils.py @@ -6,7 +6,14 @@ author: Fabrizio Grosa , CERN """ -import numpy as np +import numpy as np # pylint: disable=import-error + + +def make_list(object) -> list: + """ + Returns the object as a list if it is not a list already. + """ + return object if isinstance(object, list) else list(object) # pylint: disable=too-many-arguments @@ -46,12 +53,7 @@ def compute_crosssection( - crosssec_unc: cross-section statistical uncertainty """ - crosssection = ( - rawy - * frac - * sigma_mb - / (2 * delta_pt * delta_y * eff_times_acc * n_events * br) - ) + crosssection = rawy * frac * sigma_mb / (2 * delta_pt * delta_y * eff_times_acc * n_events * br) if method_frac == "Nb": crosssec_unc = rawy_unc / (rawy * frac) * crosssection else: @@ -68,7 +70,7 @@ def compute_fraction_fc( cross_sec_fd, raa_prompt=1.0, raa_fd=1.0, -): +) -> "tuple[list[float], list[float]]": """ Method to get fraction of prompt / FD fraction with fc method @@ -89,16 +91,13 @@ def compute_fraction_fc( - frac_fd: list of fraction of non-prompt D (central, min, max) """ - if not isinstance(cross_sec_prompt, list) and isinstance(cross_sec_prompt, float): - cross_sec_prompt = [cross_sec_prompt] - if not isinstance(cross_sec_fd, list) and isinstance(cross_sec_fd, float): - cross_sec_fd = [cross_sec_fd] - if not isinstance(raa_prompt, list) and isinstance(raa_prompt, float): - raa_prompt = [raa_prompt] - if not isinstance(raa_fd, list) and isinstance(raa_fd, float): - raa_fd = [raa_fd] + cross_sec_prompt = make_list(cross_sec_prompt) + cross_sec_fd = make_list(cross_sec_fd) + raa_prompt = make_list(raa_prompt) + raa_fd = make_list(raa_fd) - frac_prompt, frac_fd = [], [] + frac_prompt: list[float] = [] + frac_fd: list[float] = [] if acc_eff_prompt == 0: frac_fd_cent = 1.0 frac_prompt_cent = 0.0 @@ -115,37 +114,11 @@ def compute_fraction_fc( for i_sigma, (sigma_p, sigma_f) in enumerate(zip(cross_sec_prompt, cross_sec_fd)): for i_raa, (raa_p, raa_f) in enumerate(zip(raa_prompt, raa_fd)): if i_sigma == 0 and i_raa == 0: - frac_prompt_cent = 1.0 / ( - 1 + acc_eff_fd / acc_eff_prompt * sigma_f / sigma_p * raa_f / raa_p - ) - frac_fd_cent = 1.0 / ( - 1 + acc_eff_prompt / acc_eff_fd * sigma_p / sigma_f * raa_p / raa_f - ) + frac_prompt_cent = 1.0 / (1 + acc_eff_fd / acc_eff_prompt * sigma_f / sigma_p * raa_f / raa_p) + frac_fd_cent = 1.0 / (1 + acc_eff_prompt / acc_eff_fd * sigma_p / sigma_f * raa_p / raa_f) else: - frac_prompt.append( - 1.0 - / ( - 1 - + acc_eff_fd - / acc_eff_prompt - * sigma_f - / sigma_p - * raa_f - / raa_p - ) - ) - frac_fd.append( - 1.0 - / ( - 1 - + acc_eff_prompt - / acc_eff_fd - * sigma_p - / sigma_f - * raa_p - / raa_f - ) - ) + frac_prompt.append(1.0 / (1 + acc_eff_fd / acc_eff_prompt * sigma_f / sigma_p * raa_f / raa_p)) + frac_fd.append(1.0 / (1 + acc_eff_prompt / acc_eff_fd * sigma_p / sigma_f * raa_p / raa_f)) if frac_prompt and frac_fd: frac_prompt.sort() @@ -172,7 +145,7 @@ def compute_fraction_nb( sigma_mb, raa_ratio=1.0, taa=1.0, -): +) -> "list[float]": """ Method to get fraction of prompt / FD fraction with Nb method @@ -196,102 +169,39 @@ def compute_fraction_nb( - frac: list of fraction of prompt (non-prompt) D (central, min, max) """ - if not isinstance(crosssection, list) and isinstance(crosssection, float): - crosssection = [crosssection] - - if not isinstance(raa_ratio, list) and isinstance(raa_ratio, float): - raa_ratio = [raa_ratio] + crosssection = make_list(crosssection) + raa_ratio = make_list(raa_ratio) - frac = [] + frac: list[float] = [] for i_sigma, sigma in enumerate(crosssection): for i_raa_ratio, raa_rat in enumerate(raa_ratio): raa_other = 1.0 if i_sigma == 0 and i_raa_ratio == 0: if raa_rat == 1.0 and taa == 1.0: # pp - frac_cent = ( - 1 - - sigma - * delta_pt - * delta_y - * acc_eff_other - * br - * n_events - * 2 - / rawy - / sigma_mb - ) + frac_cent = 1 - sigma * delta_pt * delta_y * acc_eff_other * br * n_events * 2 / rawy / sigma_mb else: # p-Pb or Pb-Pb: iterative evaluation of Raa needed delta_raa = 1.0 while delta_raa > 1.0e-3: raw_fd = ( - taa - * raa_rat - * raa_other - * sigma - * delta_pt - * delta_y - * acc_eff_other - * br - * n_events - * 2 + taa * raa_rat * raa_other * sigma * delta_pt * delta_y * acc_eff_other * br * n_events * 2 ) frac_cent = 1 - raw_fd / rawy raa_other_old = raa_other - raa_other = ( - frac_cent - * rawy - * sigma_mb - / 2 - / acc_eff_same - / delta_pt - / delta_y - / br - / n_events - ) + raa_other = frac_cent * rawy * sigma_mb / 2 / acc_eff_same / delta_pt / delta_y / br / n_events delta_raa = abs((raa_other - raa_other_old) / raa_other) else: if raa_rat == 1.0 and taa == 1.0: # pp - frac.append( - 1 - - sigma - * delta_pt - * delta_y - * acc_eff_other - * br - * n_events - * 2 - / rawy - / sigma_mb - ) + frac.append(1 - sigma * delta_pt * delta_y * acc_eff_other * br * n_events * 2 / rawy / sigma_mb) else: # p-Pb or Pb-Pb: iterative evaluation of Raa needed delta_raa = 1.0 frac_tmp = 1.0 while delta_raa > 1.0e-3: raw_fd = ( - taa - * raa_rat - * raa_other - * sigma - * delta_pt - * delta_y - * acc_eff_other - * br - * n_events - * 2 + taa * raa_rat * raa_other * sigma * delta_pt * delta_y * acc_eff_other * br * n_events * 2 ) frac_tmp = 1 - raw_fd / rawy raa_other_old = raa_other - raa_other = ( - frac_tmp - * rawy - * sigma_mb - / 2 - / acc_eff_same - / delta_pt - / delta_y - / br - / n_events - ) + raa_other = frac_tmp * rawy * sigma_mb / 2 / acc_eff_same / delta_pt / delta_y / br / n_events delta_raa = abs((raa_other - raa_other_old) / raa_other) frac.append(frac_tmp) @@ -323,8 +233,6 @@ def get_hist_binlimits(histo): n_limits = histo.GetNbinsX() + 1 low_edge = histo.GetBinLowEdge(1) bin_width = histo.GetBinWidth(1) - bin_limits = np.array( - [low_edge + i_bin * bin_width for i_bin in range(n_limits)], "d" - ) + bin_limits = np.array([low_edge + i_bin * bin_width for i_bin in range(n_limits)], "d") return bin_limits diff --git a/PWGHF/D2H/Macros/hf_pt_spectrum.py b/PWGHF/D2H/Macros/hf_pt_spectrum.py index 620f3773c30..488a3308397 100644 --- a/PWGHF/D2H/Macros/hf_pt_spectrum.py +++ b/PWGHF/D2H/Macros/hf_pt_spectrum.py @@ -79,10 +79,7 @@ def load_inputs(input_cfg): frac_method = input_cfg["fraction"] if frac_method not in ["Nb", "fc"]: - print( - f"\033[91mERROR: method to subtract nonprompt" - f" {frac_method} not supported. Exit\033[0m" - ) + print(f"\033[91mERROR: method to subtract nonprompt" f" {frac_method} not supported. Exit\033[0m") sys.exit(5) rawy_file_name = input_cfg["rawyield"]["filename"] @@ -100,18 +97,12 @@ def load_inputs(input_cfg): infile_rawy = TFile.Open(rawy_file_name) histos["rawyields"] = infile_rawy.Get(rawy_hist_name) if not histos["rawyields"]: - print( - f"\033[91mERROR: raw-yield histo {rawy_hist_name}" - f" not found in {rawy_file_name}. Exit\033[0m" - ) + print(f"\033[91mERROR: raw-yield histo {rawy_hist_name}" f" not found in {rawy_file_name}. Exit\033[0m") sys.exit(6) histos["rawyields"].SetDirectory(0) h_events = infile_rawy.Get(norm_hist_name) if not h_events: - print( - f"\033[91mERROR: normalisation histo {norm_hist_name}" - f" not found in {rawy_file_name}. Exit\033[0m" - ) + print(f"\033[91mERROR: normalisation histo {norm_hist_name}" f" not found in {rawy_file_name}. Exit\033[0m") sys.exit(7) h_events.SetDirectory(0) infile_rawy.Close() @@ -146,14 +137,10 @@ def load_inputs(input_cfg): histos["FONLL"] = {"prompt": {}, "nonprompt": {}} infile_fonll = TFile.Open(pred_file_name) for pred in ("central", "min", "max"): - histos["FONLL"]["nonprompt"][pred] = infile_fonll.Get( - f"{fonll_hist_name[channel]}fromBpred_{pred}_corr" - ) + histos["FONLL"]["nonprompt"][pred] = infile_fonll.Get(f"{fonll_hist_name[channel]}fromBpred_{pred}_corr") histos["FONLL"]["nonprompt"][pred].SetDirectory(0) if frac_method == "fc": - histos["FONLL"]["prompt"][pred] = infile_fonll.Get( - f"{fonll_hist_name[channel]}pred_{pred}" - ) + histos["FONLL"]["prompt"][pred] = infile_fonll.Get(f"{fonll_hist_name[channel]}pred_{pred}") histos["FONLL"]["prompt"][pred].SetDirectory(0) infile_fonll.Close() @@ -163,9 +150,7 @@ def load_inputs(input_cfg): norm_db = yaml.safe_load(yml_norm_db) norm["BR"] = norm_db["BR"][channel]["value"] norm["events"] = h_events.GetBinContent(1) - norm["sigmaMB"] = ( - norm_db["sigma"]["Run2"][system][energy] if observable == "dsigmadpt" else 1.0 - ) + norm["sigmaMB"] = norm_db["sigma"]["Run2"][system][energy] if observable == "dsigmadpt" else 1.0 return histos, norm @@ -182,16 +167,14 @@ def main(): default="config_Dplus_pp5TeV.yml", help="input yaml config file", ) - parser.add_argument( - "--batch", action="store_true", default=False, help="suppress video output" - ) + parser.add_argument("--batch", action="store_true", default=False, help="suppress video output") args = parser.parse_args() if args.batch: gROOT.SetBatch(True) # final plots style settings - style_hist = TStyle('style_hist','Histo graphics style') + style_hist = TStyle("style_hist", "Histo graphics style") style_hist.SetOptStat("n") style_hist.SetMarkerColor(kAzure + 4) style_hist.SetMarkerStyle(kFullCircle) @@ -217,10 +200,7 @@ def main(): ptlims = {} for histo in ["rawyields", "acceffp", "acceffnp"]: ptlims[histo] = get_hist_binlimits(histos[histo]) - if ( - histo != "rawyields" - and not np.equal(ptlims[histo], ptlims["rawyields"]).all() - ): + if histo != "rawyields" and not np.equal(ptlims[histo], ptlims["rawyields"]).all(): print("\033[91mERROR: histo binning not consistent. Exit\033[0m") sys.exit(10) @@ -245,31 +225,22 @@ def main(): ptlims["rawyields"], ) - for i_pt, (ptmin, ptmax) in enumerate( - zip(ptlims["rawyields"][:-1], ptlims["rawyields"][1:]) - ): + for i_pt, (ptmin, ptmax) in enumerate(zip(ptlims["rawyields"][:-1], ptlims["rawyields"][1:])): pt_cent = (ptmax + ptmin) / 2 pt_delta = ptmax - ptmin rawy = histos["rawyields"].GetBinContent(i_pt + 1) rawy_unc = histos["rawyields"].GetBinError(i_pt + 1) eff_times_acc_prompt = histos["acceffp"].GetBinContent(i_pt + 1) eff_times_acc_nonprompt = histos["acceffnp"].GetBinContent(i_pt + 1) - ptmin_fonll = ( - histos["FONLL"]["nonprompt"]["central"].GetXaxis().FindBin(ptmin * 1.0001) - ) - ptmax_fonll = ( - histos["FONLL"]["nonprompt"]["central"].GetXaxis().FindBin(ptmax * 0.9999) - ) + ptmin_fonll = histos["FONLL"]["nonprompt"]["central"].GetXaxis().FindBin(ptmin * 1.0001) + ptmax_fonll = histos["FONLL"]["nonprompt"]["central"].GetXaxis().FindBin(ptmax * 0.9999) crosssec_nonprompt_fonll = [ - histos["FONLL"]["nonprompt"][pred].Integral( - ptmin_fonll, ptmax_fonll, "width" - ) - / (ptmax - ptmin) + histos["FONLL"]["nonprompt"][pred].Integral(ptmin_fonll, ptmax_fonll, "width") / (ptmax - ptmin) for pred in histos["FONLL"]["nonprompt"] ] # compute prompt fraction - frac = [0., 0., 0.] + frac = [0.0, 0.0, 0.0] if frac_method == "Nb": frac = compute_fraction_nb( # BR already included in FONLL prediction rawy, @@ -284,10 +255,7 @@ def main(): ) elif frac_method == "fc": crosssec_prompt_fonll = [ - histos["FONLL"]["prompt"][pred].Integral( - ptmin_fonll, ptmax_fonll, "width" - ) - / (ptmax - ptmin) + histos["FONLL"]["prompt"][pred].Integral(ptmin_fonll, ptmax_fonll, "width") / (ptmax - ptmin) for pred in histos["FONLL"]["prompt"] ] frac, _ = compute_fraction_fc( @@ -316,12 +284,10 @@ def main(): hptspectrum_wo_br.SetBinContent(i_pt + 1, crosssec) hptspectrum_wo_br.SetBinError(i_pt + 1, crosssec_unc) gfraction.SetPoint(i_pt, pt_cent, frac[0]) - gfraction.SetPointError( - i_pt, pt_delta / 2, pt_delta / 2, frac[0] - frac[1], frac[2] - frac[0] - ) + gfraction.SetPointError(i_pt, pt_delta / 2, pt_delta / 2, frac[0] - frac[1], frac[2] - frac[0]) c = TCanvas("c", "c", 600, 800) - c.Divide (1, 2) + c.Divide(1, 2) c.cd(1) gPad.SetLogy(True) hptspectrum.Draw() @@ -336,9 +302,7 @@ def main(): c.Print(os.path.join(output_dir, f'{cfg["output"]["filename"]}.pdf')) # save output file - output_file = TFile( - os.path.join(output_dir, f'{cfg["output"]["filename"]}.root'), "recreate" - ) + output_file = TFile(os.path.join(output_dir, f'{cfg["output"]["filename"]}.root'), "recreate") hptspectrum.Write() hptspectrum_wo_br.Write() gfraction.Write()