From 36e25b14a651414ec338ee90c1ec8c541cdce305 Mon Sep 17 00:00:00 2001 From: Fabrizio Grosa Date: Thu, 23 Jun 2022 17:02:00 +0200 Subject: [PATCH 01/10] Add first draft of python class for cut-variation method --- PWGHF/D2H/Macros/cut_variation.py | 362 ++++++++++++++++++++++++++++ PWGHF/D2H/Macros/style_formatter.py | 289 ++++++++++++++++++++++ 2 files changed, 651 insertions(+) create mode 100644 PWGHF/D2H/Macros/cut_variation.py create mode 100644 PWGHF/D2H/Macros/style_formatter.py diff --git a/PWGHF/D2H/Macros/cut_variation.py b/PWGHF/D2H/Macros/cut_variation.py new file mode 100644 index 00000000000..ca649782bd4 --- /dev/null +++ b/PWGHF/D2H/Macros/cut_variation.py @@ -0,0 +1,362 @@ +""" +Module for the (non-)prompt fraction calculation with the cut-variation method +""" + +import sys +import numpy as np +import ROOT # pylint: disable=import-error,no-name-in-module +from style_formatter import set_global_style, set_object_style + + +# pylint: disable=too-many-instance-attributes +class CutVarMinimiser: + """ + Class for the minimisation of the system of equations defined in the cut-variation method + + Parameters + ------------------------------------------------- + - raw_yields: array of floats + array of raw yields corresponding to each selection + - eff_prompt: array of floats + array of efficiencies for prompt charm hadrons corresponding to each selection + - eff_nonprompt: array of floats + array of efficiencies for non-prompt charm hadrons corresponding to each selection + """ + + def __init__( # pylint: disable=too-many-arguments + self, + raw_yields=None, + eff_prompt=None, + eff_nonprompt=None, + unc_raw_yields=None, + unc_eff_prompt=None, + unc_eff_nonprompt=None + ): + self.raw_yields = raw_yields + self.eff_prompt = eff_prompt + self.eff_nonprompt = eff_nonprompt + self.unc_raw_yields = unc_raw_yields + self.unc_eff_prompt = unc_eff_prompt + self.unc_eff_nonprompt = unc_eff_nonprompt + + self.n_sets = len(raw_yields) + + self.m_rawy = None + self.m_eff = None + self.m_cov_sets = None + self.m_corr_sets = None + self.m_weights = None + self.m_res = None + self.m_corr_yields = None + self.m_covariance = None + + self.chi_2 = 0. + self.ndf = self.n_sets - 2 + + def __check_input_consistency(self): + """ + Helper method to check self consistency of inputs + """ + + self.n_sets = len(self.raw_yields) + self.ndf = self.n_sets - 2 + + if len(self.eff_prompt) != self.n_sets or len(self.eff_nonprompt) != self.n_sets: + print("ERROR: number of raw yields and eff not consistent! Exit") + sys.exit() + + if len(self.unc_raw_yields) != self.n_sets: + print("ERROR: number of raw yield uncertainties not consistent! Exit") + sys.exit() + + if len(self.unc_eff_prompt) != self.n_sets or len(self.unc_eff_nonprompt) != self.n_sets: + print("ERROR: number of raw yield uncertainties not consistent! Exit") + sys.exit() + + def __initialise_objects(self): + """ + Helper method to initialise objects + """ + + self.m_rawy = np.zeros(shape=(self.n_sets, 1)) + self.m_eff = np.zeros(shape=(self.n_sets, 2)) + self.m_cov_sets = np.zeros(shape=(self.n_sets, self.n_sets)) + self.m_corr_sets = np.zeros(shape=(self.n_sets, self.n_sets)) + self.m_weights = np.zeros(shape=(self.n_sets, self.n_sets)) + self.m_res = np.zeros(shape=(self.n_sets, 1)) + self.m_corr_yields = np.zeros(shape=(2, 1)) + self.m_covariance = np.zeros(shape=(2, 2)) + + for i_set, (rawy, effp, effnp) in enumerate( + zip(self.raw_yields, self.eff_prompt, self.eff_nonprompt) + ): + self.m_rawy.itemset(i_set, rawy) + self.m_eff.itemset((i_set, 0), effp) + self.m_eff.itemset((i_set, 1), effnp) + + # pylint: disable=too-many-locals + def minimise_system(self, correlated=True, precision=1.e-8, max_iterations=100): + """ + Minimise the system of equations to compute the corrected yields + + Parameters + ----------------------------------------------------- + - correlated: bool + correlation between cut sets + - precision: float + target precision for minimisation procedure + - max_iterations: int + max number of iterations for minimisation procedure + """ + + self.__check_input_consistency() + self.__initialise_objects() + + self.m_rawy = np.matrix(self.m_rawy) + self.m_eff = np.matrix(self.m_eff) + m_corr_yields_old = np.zeros(shape=(2, 1)) + + for _ in range(max_iterations): + for i_row, (rw_unc_row, effp_unc_row, effnp_unc_row) in enumerate( + zip(self.unc_raw_yields, self.unc_eff_prompt, self.unc_eff_nonprompt) + ): + for i_col, (rw_unc_col, effp_unc_col, effnp_unc_col) in enumerate( + zip(self.unc_raw_yields, self.unc_eff_prompt, self.unc_eff_nonprompt)): + unc_row = np.sqrt(rw_unc_row**2 + effp_unc_row**2 * self.m_corr_yields.item(0) + ** 2 + effnp_unc_row**2 * self.m_corr_yields.item(1)**2) + unc_col = np.sqrt(rw_unc_col**2 + effp_unc_col**2 * self.m_corr_yields.item(0) + ** 2 + effnp_unc_col**2 * self.m_corr_yields.item(1)**2) + + if correlated and unc_row > 0 and unc_col > 0: + if unc_row < unc_col: + rho = unc_row / unc_col + else: + rho = unc_col / unc_row + else: + if i_row == i_col: + rho = 1. + else: + rho = 0. + cov_row_col = rho * unc_row * unc_col + self.m_cov_sets.itemset((i_row, i_col), cov_row_col) + self.m_cov_sets.itemset((i_row, i_col), rho) + + self.m_cov_sets = np.matrix(self.m_cov_sets) + self.m_weights = np.linalg.inv(np.linalg.cholesky(self.m_cov_sets)) + self.m_weights = self.m_weights.T * self.m_weights + m_eff_tr = self.m_eff.T + + self.m_covariance = (m_eff_tr * self.m_weights) * self.m_eff + self.m_covariance = np.linalg.inv(np.linalg.cholesky(self.m_covariance)) + self.m_covariance = self.m_covariance.T * self.m_covariance + + self.m_corr_yields = self.m_covariance * (m_eff_tr * self.m_weights) * self.m_rawy + self.m_res = self.m_eff * self.m_corr_yields - self.m_rawy + m_res_tr = np.transpose(self.m_res) + + rel_delta = [ + (self.m_corr_yields.item(0)-m_corr_yields_old.item(0)) / self.m_corr_yields.item(0), + (self.m_corr_yields.item(1)-m_corr_yields_old.item(1)) / self.m_corr_yields.item(1) + ] + + if rel_delta[0] < precision and rel_delta[1] < precision: + break + + m_corr_yields_old = np.copy(self.m_corr_yields) + + # chi2 + self.chi_2 = m_res_tr * self.m_weights * self.m_res + + # pylint: disable=no-member + def plot_result(self, suffix=""): + """ + Helper function to plot minimisation result as a function of cut set + + Parameters + ----------------------------------------------------- + - suffix: str + suffix to be added in the name of the output objects + + Returns + ----------------------------------------------------- + - canvas: ROOT.TCanvas + canvas with plot + - histos: dict + dictionary of ROOT.TH1F with raw yield distributions for + data, prompt, nonprompt, and the sum of prompt and nonprompt + """ + + set_global_style() + + hist_raw_yield = ROOT.TH1F( + f"hRawYieldVsCut{suffix}", + ";cut set;raw yield", + self.n_sets, + -0.5, + self.n_sets-0.5 + ) + hist_raw_yield_prompt = ROOT.TH1F( + f"hRawYieldPromptVsCut{suffix}", + ";cut set;prompt raw yield", + self.n_sets, + -0.5, + self.n_sets-0.5 + ) + hist_raw_yield_nonprompt = ROOT.TH1F( + f"hRawYieldNonPromptVsCut{suffix}", + ";cut set;non-prompt raw yield", + self.n_sets, + -0.5, + self.n_sets-0.5 + ) + hist_raw_yield_sum = ROOT.TH1F( + f"hRawYieldSumVsCut{suffix}", + ";cut set;raw yield", + self.n_sets, + -0.5, + self.n_sets-0.5 + ) + + for i_bin, (rawy, unc_rawy, effp, effnp) in enumerate( + zip(self.raw_yields, self.unc_raw_yields, self.eff_prompt, self.eff_nonprompt) + ): + hist_raw_yield.SetBinContent(i_bin + 1, rawy) + hist_raw_yield.SetBinError(i_bin + 1, unc_rawy) + + rawy_prompt = self.m_corr_yields.item(0) * effp + unc_rawy_prompt = self.m_covariance.item(0, 0) * effp + rawy_nonprompt = self.m_corr_yields.item(0) * effnp + unc_rawy_nonprompt = self.m_covariance.item(1, 1) * effnp + unc_sum = np.sqrt(unc_rawy_prompt**2 + unc_rawy_nonprompt**2 + + 2 * self.m_covariance.item(1, 0) * effp * effnp) + + hist_raw_yield_prompt.SetBinContent(i_bin + 1, rawy_prompt) + hist_raw_yield_prompt.SetBinError(i_bin + 1, unc_rawy_prompt) + hist_raw_yield_nonprompt.SetBinContent(i_bin + 1, rawy_nonprompt) + hist_raw_yield_nonprompt.SetBinError(i_bin + 1, unc_rawy_nonprompt) + hist_raw_yield_sum.SetBinContent(i_bin + 1, rawy_prompt + rawy_nonprompt) + hist_raw_yield_sum.SetBinError(i_bin + 1, unc_sum) + + set_object_style(hist_raw_yield) + set_object_style(hist_raw_yield_prompt, color=ROOT.kRed+1, alpha=0.5, fillstyle=1000) + set_object_style(hist_raw_yield_nonprompt, color=ROOT.kAzure+4, alpha=0.5, fillstyle=1000) + set_object_style(hist_raw_yield_sum, color=ROOT.kGreen+2, fillstyle=0) + + canvas = ROOT.TCanvas(f"cRawYieldVsCut{suffix}", "", 500, 500) + canvas.DrawFrame(-0.5, 0., self.n_sets - 0.5, hist_raw_yield.GetMaximum() * 1.2, + ";cut set;raw yield") + hist_raw_yield.Draw("esame") + hist_raw_yield_prompt.Draw("histsame") + hist_raw_yield_nonprompt.Draw("histsame") + hist_raw_yield_sum.Draw("histsame") + histos = { + "data": hist_raw_yield, + "prompt": hist_raw_yield_prompt, + "nonprompt": hist_raw_yield_nonprompt, + "sum": hist_raw_yield_sum, + } + + return canvas, histos + + def plot_cov_matrix(self, correlated=True, suffix=""): + """ + Helper function to plot covariance matrix + + Parameters + ----------------------------------------------------- + - correlated: bool + correlation between cut sets + - suffix: str + suffix to be added in the name of the output objects + + Returns + ----------------------------------------------------- + - canvas: ROOT.TCanvas + canvas with plot + - hist_corr_matrix: ROOT.TH2F + histogram of correlation matrix + """ + + set_global_style(palette=ROOT.kRainBow) + + hist_corr_matrix = ROOT.TH2F( + f"hCorrMatrixCutSets{suffix}", + ";cut set;cut set", + self.n_sets, + -0.5, + self.n_sets-0.5, + self.n_sets, + -0.5, + self.n_sets-0.5 + ) + for i_row, unc_row in enumerate(self.unc_raw_yields): + for i_col, unc_col in enumerate(self.unc_raw_yields): + if correlated and unc_row > 0 and unc_col > 0: + if unc_row < unc_col: + rho = unc_row / unc_col + else: + rho = unc_col / unc_row + else: + if i_row == i_col: + rho = 1. + else: + rho = 0. + hist_corr_matrix.SetBinContent(i_row, i_col, rho) + + canvas = ROOT.TCanvas(f"cCorrMatrixCutSets{suffix}", "", 500, 500) + hist_corr_matrix.Draw("colz") + + return canvas, hist_corr_matrix + + def plot_efficiencies(self, suffix=""): + """ + Helper function to plot efficiencies as a function of cut set + + Parameters + ----------------------------------------------------- + - suffix: str + suffix to be added in the name of the output objects + + Returns + ----------------------------------------------------- + - canvas: ROOT.TCanvas + canvas with plot + - histos: dict + dictionary of ROOT.TH1F with raw yield distributions for + data, prompt, nonprompt, and the sum of prompt and nonprompt + """ + + hist_eff_prompt = ROOT.TH1F( + f"hEffPromptVsCut{suffix}", + ";cut set;prompt raw yield", + self.n_sets, + -0.5, + self.n_sets-0.5 + ) + hist_eff_nonprompt = ROOT.TH1F( + f"hEffNonPromptVsCut{suffix}", + ";cut set;non-prompt raw yield", + self.n_sets, + -0.5, + self.n_sets-0.5 + ) + + for i_bin, (effp, effnp, unc_effp, unc_effnp) in enumerate( + zip(self.eff_prompt, self.eff_nonprompt, self.unc_eff_prompt, self.unc_eff_nonprompt) + ): + hist_eff_prompt.SetBinContent(i_bin + 1, effp) + hist_eff_prompt.SetBinError(i_bin + 1, unc_effp) + hist_eff_nonprompt.SetBinContent(i_bin + 1, effnp) + hist_eff_nonprompt.SetBinError(i_bin + 1, unc_effnp) + + canvas = ROOT.TCanvas(f"cEffVsCut{suffix}", "", 500, 500) + canvas.DrawFrame(-0.5, 0., self.n_sets - 0.5, hist_eff_nonprompt.GetMaximum() * 1.2, + ";cut set;efficiency") + hist_eff_prompt.Draw("histsame") + hist_eff_nonprompt.Draw("histsame") + histos = { + "prompt": hist_eff_prompt, + "nonprompt": hist_eff_nonprompt + } + + return canvas, histos diff --git a/PWGHF/D2H/Macros/style_formatter.py b/PWGHF/D2H/Macros/style_formatter.py new file mode 100644 index 00000000000..80f6a1f6ceb --- /dev/null +++ b/PWGHF/D2H/Macros/style_formatter.py @@ -0,0 +1,289 @@ +""" +Script with helper methods for style settings +""" + +import ROOT + +# pylint: disable=too-many-branches,too-many-statements, no-member + + +def set_global_style(**kwargs): + """ + Method to set global style of ROOT plots + See https://root.cern.ch/doc/master/classTStyle.html for more details + + Parameters + ------------------------------------------------- + + - padrightmargin: float + pad right margin + - padleftmargin: float + pad left margin + - padtopmargin: float + pad top margin + - padbottommargin: float + pad bottom margin + + - titlesize: float + title size for x, y, z axes + - titlesizex: float + title size for x axis + - titlesizey: float + title size for y axis + - titlesizez: float + title size for z axis + + - labelsize: float + label size for x, y, z axes + - labelsizex: float + label size for x axis + - labelsizey: float + label size for y axis + - labelsizez: float + label size for z axis + + - titleoffset: float + title offset for x, y, z axes + - titleoffsetx: float + title offset for x axis + - titleoffsety: float + title offset for y axis + - titleoffsetz: float + title offset for z axis + + - opttitle: int + title option + + - optstat: int + stats option + + - padtickx: int + pad tick option for x axis + - padticky: int + pad tick option for y axis + + - maxdigits: int + max digits for axes + + - palette: int + palette for 2D plots + """ + + # pad margins + if "padrightmargin" in kwargs: + ROOT.gStyle.SetPadRightMargin(kwargs["padrightmargin"]) + else: + ROOT.gStyle.SetPadRightMargin(0.035) + + if "padleftmargin" in kwargs: + ROOT.gStyle.SetPadLeftMargin(kwargs["padleftmargin"]) + else: + ROOT.gStyle.SetPadLeftMargin(0.12) + + if "padtopmargin" in kwargs: + ROOT.gStyle.SetPadTopMargin(kwargs["padtopmargin"]) + else: + ROOT.gStyle.SetPadTopMargin(0.035) + + if "padbottommargin" in kwargs: + ROOT.gStyle.SetPadBottomMargin(kwargs["padbottommargin"]) + else: + ROOT.gStyle.SetPadBottomMargin(0.1) + + # title sizes + if "titlesize" in kwargs: + ROOT.gStyle.SetTitleSize(kwargs["titlesize"], "xyz") + else: + ROOT.gStyle.SetTitleSize(0.050, "xyz") + + if "titlesizex" in kwargs: + ROOT.gStyle.SetTitleSize(kwargs["titlesizex"], "x") + if "titlesizey" in kwargs: + ROOT.gStyle.SetTitleSize(kwargs["titlesizex"], "y") + if "titlesizez" in kwargs: + ROOT.gStyle.SetTitleSize(kwargs["titlesizex"], "z") + + # label sizes + if "labelsize" in kwargs: + ROOT.gStyle.SetLabelSize(kwargs["labelsize"], "xyz") + else: + ROOT.gStyle.SetLabelSize(0.045, "xyz") + + if "labelsizex" in kwargs: + ROOT.gStyle.SetLabelSize(kwargs["labelsizex"], "x") + if "labelsizey" in kwargs: + ROOT.gStyle.SetLabelSize(kwargs["labelsizex"], "y") + if "labelsizez" in kwargs: + ROOT.gStyle.SetLabelSize(kwargs["labelsizex"], "z") + + # title offsets + if "titleoffset" in kwargs: + ROOT.gStyle.SetTitleOffset(kwargs["titleoffset"], "xyz") + else: + ROOT.gStyle.SetTitleOffset(1.2, "xyz") + + if "titleoffsetx" in kwargs: + ROOT.gStyle.SetTitleOffset(kwargs["titleoffsetx"], "x") + if "titleoffsety" in kwargs: + ROOT.gStyle.SetTitleOffset(kwargs["titleoffsety"], "y") + if "titleoffsetz" in kwargs: + ROOT.gStyle.SetTitleOffset(kwargs["titleoffsetz"], "z") + + # other options + if "opttitle" in kwargs: + ROOT.gStyle.SetOptTitle(kwargs["opttitle"]) + else: + ROOT.gStyle.SetOptTitle(0) + + if "optstat" in kwargs: + ROOT.gStyle.SetOptStat(kwargs["optstat"]) + else: + ROOT.gStyle.SetOptStat(0) + + if "padtickx" in kwargs: + ROOT.gStyle.SetPadTickX(kwargs["padtickx"]) + else: + ROOT.gStyle.SetPadTickX(1) + + if "padticky" in kwargs: + ROOT.gStyle.SetPadTickY(kwargs["padticky"]) + else: + ROOT.gStyle.SetPadTickY(1) + + ROOT.gStyle.SetLegendBorderSize(0) + + if "maxdigits" in kwargs: + ROOT.TGaxis.SetMaxDigits(kwargs["maxdigits"]) + + if "palette" in kwargs: + ROOT.gStyle.SetPalette(kwargs["palette"]) + + ROOT.gROOT.ForceStyle() + + +def set_object_style(obj, **kwargs): + """ + Method to set ROOT object style + See https://root.cern/doc/master/classTHistPainter.html and + https://root.cern.ch/doc/master/classTGraphPainter.html for more details + + Parameters + ------------------------------------------------- + + - obj: ROOT.TOBject + object to set style + + - linecolor: int + line color + - linealpha: float + line alpha + - linewitdh: int + line width + - linestyle: int + line style + + - markercolor: int + marker color + - markeralpha: float + marker alpha + - markerstyle: int + marker style + - markersize: int + marker size + + - fillcolor: int + fill color + - fillalpha: float + fill alpha + - fillstyle: int + fill style + + - color: int + color for line, marker, and fill + - alpha: float + alpha for line, marker, and fill + """ + + # alpha parameters + lalpha = kwargs.get("linealpha", 1) + malpha = kwargs.get("markeralpha", 1) + falpha = kwargs.get("fillalpha", 1) + if "alpha" in kwargs: + lalpha = kwargs["alpha"] + malpha = kwargs["alpha"] + falpha = kwargs["alpha"] + if "linealpha" in kwargs: + lalpha = kwargs["linealpha"] + if "markeralpha" in kwargs: + malpha = kwargs["markeralpha"] + if "fillalpha" in kwargs: + falpha = kwargs["fillalpha"] + + # line styles + if "linecolor" in kwargs: + if lalpha < 1: + obj.SetLineColorAlpha(kwargs["linecolor"], lalpha) + else: + obj.SetLineColor(kwargs["linecolor"]) + else: + if lalpha < 1: + obj.SetLineColorAlpha(ROOT.kBlack, lalpha) + else: + obj.SetLineColor(ROOT.kBlack) + + if "linewidth" in kwargs: + obj.SetLineWidth(kwargs["linewidth"]) + else: + obj.SetLineWidth(2) + + if "linestyle" in kwargs: + obj.SetLineStyle(kwargs["linestyle"]) + else: + obj.SetLineStyle(1) + + # marker styles + if "markercolor" in kwargs: + if malpha < 1: + obj.SetMarkerColorAlpha(kwargs["markercolor"], malpha) + else: + obj.SetMarkerColor(kwargs["markercolor"]) + else: + if malpha < 1: + obj.SetMarkerColorAlpha(ROOT.kBlack, malpha) + else: + obj.SetMarkerColor(ROOT.kBlack) + + if "markersize" in kwargs: + obj.SetMarkerSize(kwargs["markersize"]) + else: + obj.SetMarkerSize(1.) + + if "markerstyle" in kwargs: + obj.SetMarkerStyle(kwargs["markerstyle"]) + else: + obj.SetMarkerStyle(ROOT.kFullCircle) + + # fill styles + if "fillcolor" in kwargs: + if falpha < 1: + obj.SetFillColorAlpha(kwargs["fillcolor"], falpha) + else: + obj.SetFillColor(kwargs["fillcolor"]) + + if "fillstyle" in kwargs: + obj.SetFillStyle(kwargs["fillstyle"]) + + #global color + if "color" in kwargs: + if lalpha < 1: + obj.SetLineColorAlpha(kwargs["color"], lalpha) + else: + obj.SetLineColor(kwargs["color"]) + if malpha < 1: + obj.SetMarkerColorAlpha(kwargs["color"], malpha) + else: + obj.SetMarkerColor(kwargs["color"]) + if falpha < 1: + obj.SetFillColorAlpha(kwargs["color"], falpha) + else: + obj.SetFillColor(kwargs["color"]) From f768c5e52916234f28bd55af7807c8e704fd8885 Mon Sep 17 00:00:00 2001 From: Fabrizio Grosa Date: Fri, 24 Jun 2022 14:32:49 +0200 Subject: [PATCH 02/10] Add example script and JSON config file --- PWGHF/D2H/Macros/compute_fraction_cutvar.py | 149 ++++++++++++ PWGHF/D2H/Macros/config_cutvar_example.json | 64 +++++ PWGHF/D2H/Macros/cut_variation.py | 244 ++++++++++++++++++-- PWGHF/D2H/Macros/style_formatter.py | 4 + 4 files changed, 440 insertions(+), 21 deletions(-) create mode 100644 PWGHF/D2H/Macros/compute_fraction_cutvar.py create mode 100644 PWGHF/D2H/Macros/config_cutvar_example.json diff --git a/PWGHF/D2H/Macros/compute_fraction_cutvar.py b/PWGHF/D2H/Macros/compute_fraction_cutvar.py new file mode 100644 index 00000000000..a0408e4195b --- /dev/null +++ b/PWGHF/D2H/Macros/compute_fraction_cutvar.py @@ -0,0 +1,149 @@ +""" +Script for the (non-)prompt fraction calculation with the cut-variation method + +\author Fabrizio Grosa , CERN +\author Fabio Catalano , Politecnico and INFN Torino +\author Stefano Politanò , Politecnico and INFN Torino +""" + +import os +import argparse +import json +import numpy as np +import ROOT +from cut_variation import CutVarMinimiser +from style_formatter import set_global_style, set_object_style + +# pylint: disable=no-member,too-many-locals,too-many-statements + + +def main(config): + """ + Main function + """ + + ROOT.gROOT.SetBatch(True) + + with open(config) as fil: + cfg = json.load(fil) + + hist_rawy, hist_effp, hist_effnp = ([] for _ in range(3)) + for filename_rawy, filename_eff in zip( + cfg["rawyields"]["inputfiles"], + cfg["efficiencies"]["inputfiles"] + ): + infile_rawy = ROOT.TFile.Open(os.path.join(cfg["rawyields"]["inputdir"], filename_rawy)) + hist_rawy.append(infile_rawy.Get(cfg["rawyields"]["histoname"])) + hist_rawy[-1].SetDirectory(0) + infile_rawy.Close() + + infile_eff = ROOT.TFile.Open(os.path.join(cfg["efficiencies"]["inputdir"], filename_eff)) + hist_effp.append(infile_eff.Get(cfg["efficiencies"]["histonames"]["prompt"])) + hist_effnp.append(infile_eff.Get(cfg["efficiencies"]["histonames"]["nonprompt"])) + hist_effp[-1].SetDirectory(0) + hist_effnp[-1].SetDirectory(0) + infile_eff.Close() + + hist_corry_prompt = hist_rawy[0].Clone("hCorrYieldsPrompt") + hist_corry_nonprompt = hist_rawy[0].Clone("hCorrYieldsNonPrompt") + hist_covariance = hist_rawy[0].Clone("hCovPromptNonPrompt") + hist_corry_prompt.GetYaxis().SetTitle("corrected yields prompt") + hist_corry_nonprompt.GetYaxis().SetTitle("corrected yields non-prompt") + hist_covariance.GetYaxis().SetTitle("#sigma(prompt, non-prompt)") + set_object_style( + hist_corry_prompt, + color=ROOT.kRed+1, + fillstyle=0, + markerstyle=ROOT.kFullCircle, + linewidth=2 + ) + set_object_style( + hist_corry_nonprompt, + color=ROOT.kAzure+4, + fillstyle=0, + markerstyle=ROOT.kFullSquare + ) + set_object_style(hist_covariance) + + output = ROOT.TFile(os.path.join(cfg["output"]["directory"], cfg["output"]["file"]), "recreate") + n_sets = len(hist_rawy) + for ipt in range(hist_rawy[0].GetNbinsX()): + pt_min = hist_rawy[0].GetXaxis().GetBinLowEdge(ipt+1) + pt_max = hist_rawy[0].GetXaxis().GetBinUpEdge(ipt+1) + + rawy, effp, effnp, unc_rawy, unc_effp, unc_effnp = (np.zeros(n_sets) for _ in range(6)) + for iset, (hrawy, heffp, heffnp) in enumerate(zip(hist_rawy, hist_effp, hist_effnp)): + rawy.itemset(iset, hrawy.GetBinContent(ipt+1)) + effp.itemset(iset, heffp.GetBinContent(ipt+1)) + effnp.itemset(iset, heffnp.GetBinContent(ipt+1)) + unc_rawy.itemset(iset, hrawy.GetBinError(ipt+1)) + unc_effp.itemset(iset, heffp.GetBinError(ipt+1)) + unc_effnp.itemset(iset, heffnp.GetBinError(ipt+1)) + + minimiser = CutVarMinimiser(rawy, effp, effnp, unc_rawy, unc_effp, unc_effnp) + minimiser.minimise_system(cfg["minimisation"]["correlated"]) + + hist_corry_prompt.SetBinContent(ipt+1, minimiser.get_prompt_yield_and_error()[0]) + hist_corry_prompt.SetBinError(ipt+1, minimiser.get_prompt_yield_and_error()[1]) + hist_corry_nonprompt.SetBinContent(ipt+1, minimiser.get_nonprompt_yield_and_error()[0]) + hist_corry_nonprompt.SetBinError(ipt+1, minimiser.get_nonprompt_yield_and_error()[1]) + hist_covariance.SetBinContent(ipt+1, minimiser.get_prompt_nonprompt_cov()) + hist_covariance.SetBinError(ipt+1, 0) + + canv_rawy, histos_rawy, leg = minimiser.plot_result(f"_pt{pt_min:.0f}_{pt_max:.0f}") + output.cd() + canv_rawy.Write() + for _, hist in histos_rawy.items(): + hist.Write() + + canv_eff, histos_eff, leg = minimiser.plot_efficiencies(f"_pt{pt_min:.0f}_{pt_max:.0f}") + output.cd() + canv_eff.Write() + for _, hist in histos_eff.items(): + hist.Write() + + canv_frac, histos_frac, leg = minimiser.plot_fractions(f"_pt{pt_min:.0f}_{pt_max:.0f}") + output.cd() + canv_frac.Write() + for _, hist in histos_frac.items(): + hist.Write() + + canv_cov, histo_cov = minimiser.plot_cov_matrix(f"_pt{pt_min:.0f}_{pt_max:.0f}") + output.cd() + canv_cov.Write() + histo_cov.Write() + + output_name_rawy_pdf = f"Distr_{cfg['output']['file'].replace('.root', '.pdf')}" + output_name_eff_pdf = f"Eff_{cfg['output']['file'].replace('.root', '.pdf')}" + output_name_frac_pdf = f"Frac_{cfg['output']['file'].replace('.root', '.pdf')}" + output_name_covmat_pdf = f"CovMatrix_{cfg['output']['file'].replace('.root', '.pdf')}" + if ipt == 0: + canv_rawy.SaveAs(f"{os.path.join(cfg['output']['directory'], output_name_rawy_pdf)}[") + canv_eff.SaveAs(f"{os.path.join(cfg['output']['directory'], output_name_eff_pdf)}[") + canv_frac.SaveAs(f"{os.path.join(cfg['output']['directory'], output_name_frac_pdf)}[") + canv_cov.SaveAs(f"{os.path.join(cfg['output']['directory'], output_name_covmat_pdf)}[") + canv_rawy.SaveAs(f"{os.path.join(cfg['output']['directory'], output_name_rawy_pdf)}") + canv_eff.SaveAs(f"{os.path.join(cfg['output']['directory'], output_name_eff_pdf)}") + canv_frac.SaveAs(f"{os.path.join(cfg['output']['directory'], output_name_frac_pdf)}") + canv_cov.SaveAs(f"{os.path.join(cfg['output']['directory'], output_name_covmat_pdf)}") + if ipt == hist_rawy[0].GetNbinsX()-1: + canv_rawy.SaveAs(f"{os.path.join(cfg['output']['directory'], output_name_rawy_pdf)}]") + canv_eff.SaveAs(f"{os.path.join(cfg['output']['directory'], output_name_eff_pdf)}]") + canv_frac.SaveAs(f"{os.path.join(cfg['output']['directory'], output_name_frac_pdf)}]") + canv_cov.SaveAs(f"{os.path.join(cfg['output']['directory'], output_name_covmat_pdf)}]") + + output.cd() + hist_corry_prompt.Write() + hist_corry_nonprompt.Write() + hist_covariance.Write() + + output.Close() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Arguments") + parser.add_argument("config", metavar="text", default="config_cutvar_example.json", + help="JSON config file") + args = parser.parse_args() + + main(args.config) diff --git a/PWGHF/D2H/Macros/config_cutvar_example.json b/PWGHF/D2H/Macros/config_cutvar_example.json new file mode 100644 index 00000000000..420b558fbe4 --- /dev/null +++ b/PWGHF/D2H/Macros/config_cutvar_example.json @@ -0,0 +1,64 @@ +{ + "rawyields": { + "inputdir": "/Users/fgrosa/cernbox/alice_work/analyses/np_Dplus_vsmult_pp13TeV/outputs/rawyields/MB_INELgtr0", + "inputfiles": [ + "RawYieldsDplus_outFD_pos10.root", + "RawYieldsDplus_outFD_pos11.root", + "RawYieldsDplus_outFD_pos12.root", + "RawYieldsDplus_outFD_pos13.root", + "RawYieldsDplus_outFD_pos14.root", + "RawYieldsDplus_outFD_pos15.root", + "RawYieldsDplus_outFD_pos16.root", + "RawYieldsDplus_outFD_pos17.root", + "RawYieldsDplus_outFD_pos18.root", + "RawYieldsDplus_outFD_pos19.root", + "RawYieldsDplus_outFD_pos20.root", + "RawYieldsDplus_outFD_pos21.root", + "RawYieldsDplus_outFD_pos22.root", + "RawYieldsDplus_outFD_pos23.root", + "RawYieldsDplus_outFD_pos24.root", + "RawYieldsDplus_outFD_pos25.root", + "RawYieldsDplus_outFD_pos26.root", + "RawYieldsDplus_outFD_pos27.root", + "RawYieldsDplus_outFD_pos28.root", + "RawYieldsDplus_outFD_pos29.root" + ], + "histoname": "hRawYields" + }, + "efficiencies": { + "inputdir": "/Users/fgrosa/cernbox/alice_work/analyses/np_Dplus_vsmult_pp13TeV/outputs/efficiencies/MB_INELgtr0", + "inputfiles": [ + "Eff_times_Acc_Dplus_outFD_pos10.root", + "Eff_times_Acc_Dplus_outFD_pos11.root", + "Eff_times_Acc_Dplus_outFD_pos12.root", + "Eff_times_Acc_Dplus_outFD_pos13.root", + "Eff_times_Acc_Dplus_outFD_pos14.root", + "Eff_times_Acc_Dplus_outFD_pos15.root", + "Eff_times_Acc_Dplus_outFD_pos16.root", + "Eff_times_Acc_Dplus_outFD_pos17.root", + "Eff_times_Acc_Dplus_outFD_pos18.root", + "Eff_times_Acc_Dplus_outFD_pos19.root", + "Eff_times_Acc_Dplus_outFD_pos20.root", + "Eff_times_Acc_Dplus_outFD_pos21.root", + "Eff_times_Acc_Dplus_outFD_pos22.root", + "Eff_times_Acc_Dplus_outFD_pos23.root", + "Eff_times_Acc_Dplus_outFD_pos24.root", + "Eff_times_Acc_Dplus_outFD_pos25.root", + "Eff_times_Acc_Dplus_outFD_pos26.root", + "Eff_times_Acc_Dplus_outFD_pos27.root", + "Eff_times_Acc_Dplus_outFD_pos28.root", + "Eff_times_Acc_Dplus_outFD_pos29.root" + ], + "histonames": { + "prompt": "hAccEffPrompt", + "nonprompt": "hAccEffFD" + } + }, + "minimisation": { + "correlated": "true" + }, + "output": { + "directory": ".", + "file": "CutVarDplus_pp13TeV_MB.root" + } +} \ No newline at end of file diff --git a/PWGHF/D2H/Macros/cut_variation.py b/PWGHF/D2H/Macros/cut_variation.py index ca649782bd4..cdbf315ffa6 100644 --- a/PWGHF/D2H/Macros/cut_variation.py +++ b/PWGHF/D2H/Macros/cut_variation.py @@ -1,10 +1,14 @@ """ Module for the (non-)prompt fraction calculation with the cut-variation method + +\author Fabrizio Grosa , CERN +\author Fabio Catalano , Politecnico and INFN Torino +\author Stefano Politanò , Politecnico and INFN Torino """ import sys import numpy as np -import ROOT # pylint: disable=import-error,no-name-in-module +import ROOT from style_formatter import set_global_style, set_object_style @@ -21,6 +25,14 @@ class CutVarMinimiser: array of efficiencies for prompt charm hadrons corresponding to each selection - eff_nonprompt: array of floats array of efficiencies for non-prompt charm hadrons corresponding to each selection + - unc_raw_yields: array of floats + array of raw yield uncertainties corresponding to each selection + - unc_eff_prompt: array of floats + array of efficiency uncertainties for prompt charm hadrons + corresponding to each selection + - unc_eff_nonprompt: array of floats + array of efficiency uncertainties for non-prompt charm hadrons + corresponding to each selection """ def __init__( # pylint: disable=too-many-arguments @@ -39,6 +51,11 @@ def __init__( # pylint: disable=too-many-arguments self.unc_eff_prompt = unc_eff_prompt self.unc_eff_nonprompt = unc_eff_nonprompt + self.frac_prompt = None + self.frac_nonprompt = None + self.unc_frac_prompt = None + self.unc_frac_nonprompt = None + self.n_sets = len(raw_yields) self.m_rawy = None @@ -87,6 +104,11 @@ def __initialise_objects(self): self.m_corr_yields = np.zeros(shape=(2, 1)) self.m_covariance = np.zeros(shape=(2, 2)) + self.frac_prompt = np.zeros(shape=self.n_sets) + self.frac_nonprompt = np.zeros(shape=self.n_sets) + self.unc_frac_prompt = np.zeros(shape=self.n_sets) + self.unc_frac_nonprompt = np.zeros(shape=self.n_sets) + for i_set, (rawy, effp, effnp) in enumerate( zip(self.raw_yields, self.eff_prompt, self.eff_nonprompt) ): @@ -139,7 +161,6 @@ def minimise_system(self, correlated=True, precision=1.e-8, max_iterations=100): rho = 0. cov_row_col = rho * unc_row * unc_col self.m_cov_sets.itemset((i_row, i_col), cov_row_col) - self.m_cov_sets.itemset((i_row, i_col), rho) self.m_cov_sets = np.matrix(self.m_cov_sets) self.m_weights = np.linalg.inv(np.linalg.cholesky(self.m_cov_sets)) @@ -165,7 +186,55 @@ def minimise_system(self, correlated=True, precision=1.e-8, max_iterations=100): m_corr_yields_old = np.copy(self.m_corr_yields) # chi2 - self.chi_2 = m_res_tr * self.m_weights * self.m_res + self.chi_2 = np.float(m_res_tr * self.m_weights * self.m_res) + + # fraction + for i_set, (effp, effnp) in enumerate(zip(self.eff_prompt, self.eff_nonprompt)): + rawyp = effp * self.m_corr_yields.item(0) + rawynp = effnp * self.m_corr_yields.item(1) + der_fp_p = (effp * (rawyp + rawynp) - effp**2 * self.m_corr_yields.item(0)) / (rawyp + rawynp)**2 + der_fp_np = -effp * effnp * self.m_corr_yields.item(0) / (rawyp + rawynp)**2 + der_fnp_np = (effnp * (rawyp + rawynp) - effnp**2 * self.m_corr_yields.item(1)) / (rawyp + rawynp)**2 + der_fnp_p = -effp * effnp * self.m_corr_yields.item(1) / (rawyp + rawynp)**2 + + unc_fp = np.sqrt(der_fp_p**2 * self.m_covariance.item(0, 0) + + der_fp_np**2 * self.m_covariance.item(1, 1) + + 2 * der_fp_p * der_fp_np * self.m_covariance.item(1, 0)) + unc_fnp = np.sqrt(der_fnp_p**2 * self.m_covariance.item(0, 0) + + der_fnp_np**2 * self.m_covariance.item(1, 1) + + 2 * der_fnp_p * der_fnp_np * self.m_covariance.item(1, 0)) + self.frac_prompt.itemset(i_set, rawyp / (rawyp + rawynp)) + self.frac_nonprompt.itemset(i_set, rawynp / (rawyp + rawynp)) + self.unc_frac_prompt.itemset(i_set, unc_fp) + self.unc_frac_nonprompt.itemset(i_set, unc_fnp) + + def get_red_chi2(self): + """ + Helper function to get reduced chi2 + """ + + return self.chi_2 / self.ndf + + def get_prompt_yield_and_error(self): + """ + Helper function to get prompt corrected yield and error + """ + + return self.m_corr_yields.item(0), np.sqrt(self.m_covariance.item(0, 0)) + + def get_nonprompt_yield_and_error(self): + """ + Helper function to get non-prompt corrected yield and error + """ + + return self.m_corr_yields.item(1), np.sqrt(self.m_covariance.item(1, 1)) + + def get_prompt_nonprompt_cov(self): + """ + Helper function to get covariance between prompt and non-prompt corrected yields + """ + + return self.m_covariance.item(1, 0) # pylint: disable=no-member def plot_result(self, suffix=""): @@ -184,9 +253,11 @@ def plot_result(self, suffix=""): - histos: dict dictionary of ROOT.TH1F with raw yield distributions for data, prompt, nonprompt, and the sum of prompt and nonprompt + - leg: ROOT.TLegend + needed otherwise it is destroyed """ - set_global_style() + set_global_style(padleftmargin=0.16, padbottommargin=0.12, titleoffsety=1.6) hist_raw_yield = ROOT.TH1F( f"hRawYieldVsCut{suffix}", @@ -224,9 +295,9 @@ def plot_result(self, suffix=""): hist_raw_yield.SetBinError(i_bin + 1, unc_rawy) rawy_prompt = self.m_corr_yields.item(0) * effp - unc_rawy_prompt = self.m_covariance.item(0, 0) * effp - rawy_nonprompt = self.m_corr_yields.item(0) * effnp - unc_rawy_nonprompt = self.m_covariance.item(1, 1) * effnp + unc_rawy_prompt = np.sqrt(self.m_covariance.item(0, 0)) * effp + rawy_nonprompt = self.m_corr_yields.item(1) * effnp + unc_rawy_nonprompt = np.sqrt(self.m_covariance.item(1, 1)) * effnp unc_sum = np.sqrt(unc_rawy_prompt**2 + unc_rawy_nonprompt**2 + 2 * self.m_covariance.item(1, 0) * effp * effnp) @@ -238,17 +309,30 @@ def plot_result(self, suffix=""): hist_raw_yield_sum.SetBinError(i_bin + 1, unc_sum) set_object_style(hist_raw_yield) - set_object_style(hist_raw_yield_prompt, color=ROOT.kRed+1, alpha=0.5, fillstyle=1000) - set_object_style(hist_raw_yield_nonprompt, color=ROOT.kAzure+4, alpha=0.5, fillstyle=1000) + set_object_style(hist_raw_yield_prompt, color=ROOT.kRed+1, fillstyle=3145) + set_object_style(hist_raw_yield_nonprompt, color=ROOT.kAzure+4, fillstyle=3154) set_object_style(hist_raw_yield_sum, color=ROOT.kGreen+2, fillstyle=0) canvas = ROOT.TCanvas(f"cRawYieldVsCut{suffix}", "", 500, 500) canvas.DrawFrame(-0.5, 0., self.n_sets - 0.5, hist_raw_yield.GetMaximum() * 1.2, ";cut set;raw yield") + leg = ROOT.TLegend(0.6, 0.65, 0.8, 0.9) + leg.SetBorderSize(0) + leg.SetFillStyle(0) + leg.SetTextSize(0.04) + leg.SetHeader(f"#chi^{{2}}/#it{{ndf}} = {self.chi_2:.2f}/{self.ndf}") + leg.AddEntry(hist_raw_yield, "data", "p") + leg.AddEntry(hist_raw_yield_prompt, "prompt", "f") + leg.AddEntry(hist_raw_yield_nonprompt, "non-prompt", "f") + leg.AddEntry(hist_raw_yield_sum, "total", "l") + leg.Draw() hist_raw_yield.Draw("esame") hist_raw_yield_prompt.Draw("histsame") hist_raw_yield_nonprompt.Draw("histsame") hist_raw_yield_sum.Draw("histsame") + canvas.Modified() + canvas.Update() + histos = { "data": hist_raw_yield, "prompt": hist_raw_yield_prompt, @@ -256,7 +340,7 @@ def plot_result(self, suffix=""): "sum": hist_raw_yield_sum, } - return canvas, histos + return canvas, histos, leg def plot_cov_matrix(self, correlated=True, suffix=""): """ @@ -277,7 +361,12 @@ def plot_cov_matrix(self, correlated=True, suffix=""): histogram of correlation matrix """ - set_global_style(palette=ROOT.kRainBow) + set_global_style( + padleftmargin=0.14, + padbottommargin=0.12, + padrightmargin=0.12, + palette=ROOT.kRainBow + ) hist_corr_matrix = ROOT.TH2F( f"hCorrMatrixCutSets{suffix}", @@ -301,10 +390,12 @@ def plot_cov_matrix(self, correlated=True, suffix=""): rho = 1. else: rho = 0. - hist_corr_matrix.SetBinContent(i_row, i_col, rho) + hist_corr_matrix.SetBinContent(i_row+1, i_col+1, rho) canvas = ROOT.TCanvas(f"cCorrMatrixCutSets{suffix}", "", 500, 500) hist_corr_matrix.Draw("colz") + canvas.Modified() + canvas.Update() return canvas, hist_corr_matrix @@ -322,20 +413,23 @@ def plot_efficiencies(self, suffix=""): - canvas: ROOT.TCanvas canvas with plot - histos: dict - dictionary of ROOT.TH1F with raw yield distributions for - data, prompt, nonprompt, and the sum of prompt and nonprompt + dictionary of ROOT.TH1F with efficiencies for prompt and nonprompt + - leg: ROOT.TLegend + needed otherwise it is destroyed """ + set_global_style(padleftmargin=0.14, padbottommargin=0.12, titleoffset=1.2) + hist_eff_prompt = ROOT.TH1F( f"hEffPromptVsCut{suffix}", - ";cut set;prompt raw yield", + ";cut set;prompt acceptance #times efficiency", self.n_sets, -0.5, self.n_sets-0.5 ) hist_eff_nonprompt = ROOT.TH1F( f"hEffNonPromptVsCut{suffix}", - ";cut set;non-prompt raw yield", + ";cut set;non-prompt acceptance #times efficiency", self.n_sets, -0.5, self.n_sets-0.5 @@ -349,14 +443,122 @@ def plot_efficiencies(self, suffix=""): hist_eff_nonprompt.SetBinContent(i_bin + 1, effnp) hist_eff_nonprompt.SetBinError(i_bin + 1, unc_effnp) + set_object_style( + hist_eff_prompt, + color=ROOT.kRed+1, + fillstyle=0, + markerstyle=ROOT.kFullCircle, + linewidth=2 + ) + set_object_style( + hist_eff_nonprompt, + color=ROOT.kAzure+4, + fillstyle=0, + markerstyle=ROOT.kFullSquare + ) + canvas = ROOT.TCanvas(f"cEffVsCut{suffix}", "", 500, 500) - canvas.DrawFrame(-0.5, 0., self.n_sets - 0.5, hist_eff_nonprompt.GetMaximum() * 1.2, - ";cut set;efficiency") - hist_eff_prompt.Draw("histsame") - hist_eff_nonprompt.Draw("histsame") + canvas.DrawFrame(-0.5, 1.e-5, self.n_sets - 0.5, 1., + ";cut set;acceptance #times efficiency") + canvas.SetLogy() + hist_eff_prompt.Draw("esame") + hist_eff_nonprompt.Draw("esame") histos = { "prompt": hist_eff_prompt, "nonprompt": hist_eff_nonprompt } + leg = ROOT.TLegend(0.2, 0.2, 0.4, 0.3) + leg.SetBorderSize(0) + leg.SetFillStyle(0) + leg.SetTextSize(0.04) + leg.AddEntry(hist_eff_prompt, "prompt", "pl") + leg.AddEntry(hist_eff_nonprompt, "non-prompt", "pl") + leg.Draw() + canvas.Modified() + canvas.Update() + + return canvas, histos, leg + + def plot_fractions(self, suffix=""): + """ + Helper function to plot fractions as a function of cut set - return canvas, histos + Parameters + ----------------------------------------------------- + - suffix: str + suffix to be added in the name of the output objects + + Returns + ----------------------------------------------------- + - canvas: ROOT.TCanvas + canvas with plot + - histos: dict + dictionary of ROOT.TH1F with fractions for prompt and nonprompt + - leg: ROOT.TLegend + needed otherwise it is destroyed + """ + + set_global_style(padleftmargin=0.14, padbottommargin=0.12, titleoffset=1.2) + + hist_f_prompt = ROOT.TH1F( + f"hFracPromptVsCut{suffix}", + ";cut set;#it{f}_{prompt}", + self.n_sets, + -0.5, + self.n_sets-0.5 + ) + hist_f_nonprompt = ROOT.TH1F( + f"hFracNonPromptVsCut{suffix}", + ";cut set;#it{f}_{non-prompt}", + self.n_sets, + -0.5, + self.n_sets-0.5 + ) + + for i_bin, (fracp, fracnp, unc_fp, unc_fnp) in enumerate( + zip( + self.frac_prompt, + self.frac_nonprompt, + self.unc_frac_prompt, + self.unc_frac_nonprompt + ) + ): + hist_f_prompt.SetBinContent(i_bin + 1, fracp) + hist_f_prompt.SetBinError(i_bin + 1, unc_fp) + hist_f_nonprompt.SetBinContent(i_bin + 1, fracnp) + hist_f_nonprompt.SetBinError(i_bin + 1, unc_fnp) + + set_object_style( + hist_f_prompt, + color=ROOT.kRed+1, + fillstyle=0, + markerstyle=ROOT.kFullCircle, + linewidth=2 + ) + set_object_style( + hist_f_nonprompt, + color=ROOT.kAzure+4, + fillstyle=0, + markerstyle=ROOT.kFullSquare + ) + + canvas = ROOT.TCanvas(f"cFracVsCut{suffix}", "", 500, 500) + canvas.DrawFrame(-0.5, 0., self.n_sets - 0.5, 1., + ";cut set;fraction") + hist_f_prompt.Draw("esame") + hist_f_nonprompt.Draw("esame") + histos = { + "prompt": hist_f_prompt, + "nonprompt": hist_f_nonprompt + } + leg = ROOT.TLegend(0.2, 0.8, 0.4, 0.9) + leg.SetBorderSize(0) + leg.SetFillStyle(0) + leg.SetTextSize(0.04) + leg.AddEntry(hist_f_prompt, "prompt", "pl") + leg.AddEntry(hist_f_nonprompt, "non-prompt", "pl") + leg.Draw() + canvas.Modified() + canvas.Update() + + return canvas, histos, leg diff --git a/PWGHF/D2H/Macros/style_formatter.py b/PWGHF/D2H/Macros/style_formatter.py index 80f6a1f6ceb..d8de163775f 100644 --- a/PWGHF/D2H/Macros/style_formatter.py +++ b/PWGHF/D2H/Macros/style_formatter.py @@ -1,5 +1,9 @@ """ Script with helper methods for style settings + +\author Fabrizio Grosa , CERN +\author Fabio Catalano , Politecnico and INFN Torino +\author Stefano Politanò , Politecnico and INFN Torino """ import ROOT From 509077fb09aa6c3f3d1c78c64f39c5ba54459e98 Mon Sep 17 00:00:00 2001 From: Fabrizio Grosa Date: Fri, 24 Jun 2022 14:40:45 +0200 Subject: [PATCH 03/10] Add README --- PWGHF/D2H/Macros/README.md | 7 +++++++ 1 file changed, 7 insertions(+) create mode 100644 PWGHF/D2H/Macros/README.md diff --git a/PWGHF/D2H/Macros/README.md b/PWGHF/D2H/Macros/README.md new file mode 100644 index 00000000000..de44037e6b9 --- /dev/null +++ b/PWGHF/D2H/Macros/README.md @@ -0,0 +1,7 @@ +# Cut-variation method +In order to compute the (non-)prompt fraction of charm hadrons with the cut-variation method, the [`compute_fraction_cutvar.py`](https://github.com/AliceO2Group/O2Physics/blob/master/PWGHF/D2H/Macros/compute_fraction_cutvar.py) script can be used. +It requires as inputs `.root` files with transverse-momentum differential raw yields and efficiencies of prompt and non-prompt production stored in `ROOT::TH1F` objects. These have to be passed to the script via a `JSON` config file, such as [`config_cutvar_example.json`](https://github.com/AliceO2Group/O2Physics/blob/master/PWGHF/D2H/Macros/config_cutvar_example.json). The script has to be executed via +```python +python3 compute_fraction_cutvar.py config_cutvar_example.json +``` +It produces a `.root` file with output histograms and four `.pdf` files with fitted distributions, efficiencies, covariance matrices, and resulting fractions. From 29d47cfffc1fb2ac95d4a6e28a960bb5c92cc2a8 Mon Sep 17 00:00:00 2001 From: Fabrizio Grosa Date: Fri, 24 Jun 2022 14:42:29 +0200 Subject: [PATCH 04/10] Add authors --- PWGHF/D2H/Macros/compute_fraction_cutvar.py | 1 + PWGHF/D2H/Macros/cut_variation.py | 1 + PWGHF/D2H/Macros/style_formatter.py | 1 + 3 files changed, 3 insertions(+) diff --git a/PWGHF/D2H/Macros/compute_fraction_cutvar.py b/PWGHF/D2H/Macros/compute_fraction_cutvar.py index a0408e4195b..56400394a2f 100644 --- a/PWGHF/D2H/Macros/compute_fraction_cutvar.py +++ b/PWGHF/D2H/Macros/compute_fraction_cutvar.py @@ -4,6 +4,7 @@ \author Fabrizio Grosa , CERN \author Fabio Catalano , Politecnico and INFN Torino \author Stefano Politanò , Politecnico and INFN Torino +\author Daniel Battistini , TUM """ import os diff --git a/PWGHF/D2H/Macros/cut_variation.py b/PWGHF/D2H/Macros/cut_variation.py index cdbf315ffa6..2ab1aeb52c9 100644 --- a/PWGHF/D2H/Macros/cut_variation.py +++ b/PWGHF/D2H/Macros/cut_variation.py @@ -4,6 +4,7 @@ \author Fabrizio Grosa , CERN \author Fabio Catalano , Politecnico and INFN Torino \author Stefano Politanò , Politecnico and INFN Torino +\author Daniel Battistini , TUM """ import sys diff --git a/PWGHF/D2H/Macros/style_formatter.py b/PWGHF/D2H/Macros/style_formatter.py index d8de163775f..0ad3cf6d498 100644 --- a/PWGHF/D2H/Macros/style_formatter.py +++ b/PWGHF/D2H/Macros/style_formatter.py @@ -4,6 +4,7 @@ \author Fabrizio Grosa , CERN \author Fabio Catalano , Politecnico and INFN Torino \author Stefano Politanò , Politecnico and INFN Torino +\author Daniel Battistini , TUM """ import ROOT From 11bdf178ea5732a252d7247cebd619a6428fc200 Mon Sep 17 00:00:00 2001 From: Fabrizio Grosa Date: Fri, 24 Jun 2022 14:53:28 +0200 Subject: [PATCH 05/10] Fix typo --- PWGHF/D2H/Macros/compute_fraction_cutvar.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/PWGHF/D2H/Macros/compute_fraction_cutvar.py b/PWGHF/D2H/Macros/compute_fraction_cutvar.py index 56400394a2f..a9fc209d57e 100644 --- a/PWGHF/D2H/Macros/compute_fraction_cutvar.py +++ b/PWGHF/D2H/Macros/compute_fraction_cutvar.py @@ -91,19 +91,19 @@ def main(config): hist_covariance.SetBinContent(ipt+1, minimiser.get_prompt_nonprompt_cov()) hist_covariance.SetBinError(ipt+1, 0) - canv_rawy, histos_rawy, leg = minimiser.plot_result(f"_pt{pt_min:.0f}_{pt_max:.0f}") + canv_rawy, histos_rawy, leg_r = minimiser.plot_result(f"_pt{pt_min:.0f}_{pt_max:.0f}") output.cd() canv_rawy.Write() for _, hist in histos_rawy.items(): hist.Write() - canv_eff, histos_eff, leg = minimiser.plot_efficiencies(f"_pt{pt_min:.0f}_{pt_max:.0f}") + canv_eff, histos_eff, leg_e = minimiser.plot_efficiencies(f"_pt{pt_min:.0f}_{pt_max:.0f}") output.cd() canv_eff.Write() for _, hist in histos_eff.items(): hist.Write() - canv_frac, histos_frac, leg = minimiser.plot_fractions(f"_pt{pt_min:.0f}_{pt_max:.0f}") + canv_frac, histos_frac, leg_f = minimiser.plot_fractions(f"_pt{pt_min:.0f}_{pt_max:.0f}") output.cd() canv_frac.Write() for _, hist in histos_frac.items(): From f8f1cb9b30addd82a736bff0ed6f84763382c952 Mon Sep 17 00:00:00 2001 From: Fabrizio Grosa Date: Sun, 26 Jun 2022 19:09:42 +0200 Subject: [PATCH 06/10] Temporarily remove links --- PWGHF/D2H/Macros/README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/PWGHF/D2H/Macros/README.md b/PWGHF/D2H/Macros/README.md index de44037e6b9..03c9c9cb4a2 100644 --- a/PWGHF/D2H/Macros/README.md +++ b/PWGHF/D2H/Macros/README.md @@ -1,6 +1,6 @@ # Cut-variation method -In order to compute the (non-)prompt fraction of charm hadrons with the cut-variation method, the [`compute_fraction_cutvar.py`](https://github.com/AliceO2Group/O2Physics/blob/master/PWGHF/D2H/Macros/compute_fraction_cutvar.py) script can be used. -It requires as inputs `.root` files with transverse-momentum differential raw yields and efficiencies of prompt and non-prompt production stored in `ROOT::TH1F` objects. These have to be passed to the script via a `JSON` config file, such as [`config_cutvar_example.json`](https://github.com/AliceO2Group/O2Physics/blob/master/PWGHF/D2H/Macros/config_cutvar_example.json). The script has to be executed via +In order to compute the (non-)prompt fraction of charm hadrons with the cut-variation method, the `compute_fraction_cutvar.py` script can be used. +It requires as inputs `.root` files with transverse-momentum differential raw yields and efficiencies of prompt and non-prompt production stored in `ROOT::TH1F` objects. These have to be passed to the script via a `JSON` config file, such as `config_cutvar_example.json`. The script has to be executed via ```python python3 compute_fraction_cutvar.py config_cutvar_example.json ``` From bc1849134b60324a3f533c39ab45070eb62ebc86 Mon Sep 17 00:00:00 2001 From: Fabrizio Grosa Date: Sun, 26 Jun 2022 19:11:50 +0200 Subject: [PATCH 07/10] Make flake8 happy --- PWGHF/D2H/Macros/compute_fraction_cutvar.py | 2 +- PWGHF/D2H/Macros/style_formatter.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/PWGHF/D2H/Macros/compute_fraction_cutvar.py b/PWGHF/D2H/Macros/compute_fraction_cutvar.py index a9fc209d57e..772c48f0126 100644 --- a/PWGHF/D2H/Macros/compute_fraction_cutvar.py +++ b/PWGHF/D2H/Macros/compute_fraction_cutvar.py @@ -13,7 +13,7 @@ import numpy as np import ROOT from cut_variation import CutVarMinimiser -from style_formatter import set_global_style, set_object_style +from style_formatter import set_object_style # pylint: disable=no-member,too-many-locals,too-many-statements diff --git a/PWGHF/D2H/Macros/style_formatter.py b/PWGHF/D2H/Macros/style_formatter.py index 0ad3cf6d498..bb3d98d9640 100644 --- a/PWGHF/D2H/Macros/style_formatter.py +++ b/PWGHF/D2H/Macros/style_formatter.py @@ -278,7 +278,7 @@ def set_object_style(obj, **kwargs): if "fillstyle" in kwargs: obj.SetFillStyle(kwargs["fillstyle"]) - #global color + # global color if "color" in kwargs: if lalpha < 1: obj.SetLineColorAlpha(kwargs["color"], lalpha) From ced825ce855463825d6f68df162036bdc45805e7 Mon Sep 17 00:00:00 2001 From: Fabrizio Grosa Date: Sun, 26 Jun 2022 19:13:08 +0200 Subject: [PATCH 08/10] Fix section title --- PWGHF/D2H/Macros/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/PWGHF/D2H/Macros/README.md b/PWGHF/D2H/Macros/README.md index 03c9c9cb4a2..e8f52100595 100644 --- a/PWGHF/D2H/Macros/README.md +++ b/PWGHF/D2H/Macros/README.md @@ -1,4 +1,4 @@ -# Cut-variation method +# (non-)prompt fraction via cut-variation method In order to compute the (non-)prompt fraction of charm hadrons with the cut-variation method, the `compute_fraction_cutvar.py` script can be used. It requires as inputs `.root` files with transverse-momentum differential raw yields and efficiencies of prompt and non-prompt production stored in `ROOT::TH1F` objects. These have to be passed to the script via a `JSON` config file, such as `config_cutvar_example.json`. The script has to be executed via ```python From b7765760525ef28850859e44d2d8b08655d98590 Mon Sep 17 00:00:00 2001 From: Fabrizio Grosa Date: Fri, 1 Jul 2022 16:23:16 +0200 Subject: [PATCH 09/10] Remove README --- PWGHF/D2H/Macros/README.md | 7 ------- 1 file changed, 7 deletions(-) delete mode 100644 PWGHF/D2H/Macros/README.md diff --git a/PWGHF/D2H/Macros/README.md b/PWGHF/D2H/Macros/README.md deleted file mode 100644 index e8f52100595..00000000000 --- a/PWGHF/D2H/Macros/README.md +++ /dev/null @@ -1,7 +0,0 @@ -# (non-)prompt fraction via cut-variation method -In order to compute the (non-)prompt fraction of charm hadrons with the cut-variation method, the `compute_fraction_cutvar.py` script can be used. -It requires as inputs `.root` files with transverse-momentum differential raw yields and efficiencies of prompt and non-prompt production stored in `ROOT::TH1F` objects. These have to be passed to the script via a `JSON` config file, such as `config_cutvar_example.json`. The script has to be executed via -```python -python3 compute_fraction_cutvar.py config_cutvar_example.json -``` -It produces a `.root` file with output histograms and four `.pdf` files with fitted distributions, efficiencies, covariance matrices, and resulting fractions. From d0decada72eafd4c07efc0f4eb103bae6aff62d5 Mon Sep 17 00:00:00 2001 From: ALICE Action Bot Date: Fri, 1 Jul 2022 14:26:42 +0000 Subject: [PATCH 10/10] MegaLinter fixes --- PWGHF/D2H/Macros/compute_fraction_cutvar.py | 152 +++++++++++----- PWGHF/D2H/Macros/cut_variation.py | 190 ++++++++++++-------- PWGHF/D2H/Macros/style_formatter.py | 2 +- 3 files changed, 225 insertions(+), 119 deletions(-) diff --git a/PWGHF/D2H/Macros/compute_fraction_cutvar.py b/PWGHF/D2H/Macros/compute_fraction_cutvar.py index 772c48f0126..c9b5cf654cd 100644 --- a/PWGHF/D2H/Macros/compute_fraction_cutvar.py +++ b/PWGHF/D2H/Macros/compute_fraction_cutvar.py @@ -7,9 +7,10 @@ \author Daniel Battistini , TUM """ -import os import argparse import json +import os + import numpy as np import ROOT from cut_variation import CutVarMinimiser @@ -30,17 +31,22 @@ def main(config): hist_rawy, hist_effp, hist_effnp = ([] for _ in range(3)) for filename_rawy, filename_eff in zip( - cfg["rawyields"]["inputfiles"], - cfg["efficiencies"]["inputfiles"] + cfg["rawyields"]["inputfiles"], cfg["efficiencies"]["inputfiles"] ): - infile_rawy = ROOT.TFile.Open(os.path.join(cfg["rawyields"]["inputdir"], filename_rawy)) + infile_rawy = ROOT.TFile.Open( + os.path.join(cfg["rawyields"]["inputdir"], filename_rawy) + ) hist_rawy.append(infile_rawy.Get(cfg["rawyields"]["histoname"])) hist_rawy[-1].SetDirectory(0) infile_rawy.Close() - infile_eff = ROOT.TFile.Open(os.path.join(cfg["efficiencies"]["inputdir"], filename_eff)) + infile_eff = ROOT.TFile.Open( + os.path.join(cfg["efficiencies"]["inputdir"], filename_eff) + ) hist_effp.append(infile_eff.Get(cfg["efficiencies"]["histonames"]["prompt"])) - hist_effnp.append(infile_eff.Get(cfg["efficiencies"]["histonames"]["nonprompt"])) + hist_effnp.append( + infile_eff.Get(cfg["efficiencies"]["histonames"]["nonprompt"]) + ) hist_effp[-1].SetDirectory(0) hist_effnp[-1].SetDirectory(0) infile_eff.Close() @@ -53,57 +59,77 @@ def main(config): hist_covariance.GetYaxis().SetTitle("#sigma(prompt, non-prompt)") set_object_style( hist_corry_prompt, - color=ROOT.kRed+1, + color=ROOT.kRed + 1, fillstyle=0, markerstyle=ROOT.kFullCircle, - linewidth=2 + linewidth=2, ) set_object_style( hist_corry_nonprompt, - color=ROOT.kAzure+4, + color=ROOT.kAzure + 4, fillstyle=0, - markerstyle=ROOT.kFullSquare + markerstyle=ROOT.kFullSquare, ) set_object_style(hist_covariance) - output = ROOT.TFile(os.path.join(cfg["output"]["directory"], cfg["output"]["file"]), "recreate") + output = ROOT.TFile( + os.path.join(cfg["output"]["directory"], cfg["output"]["file"]), "recreate" + ) n_sets = len(hist_rawy) for ipt in range(hist_rawy[0].GetNbinsX()): - pt_min = hist_rawy[0].GetXaxis().GetBinLowEdge(ipt+1) - pt_max = hist_rawy[0].GetXaxis().GetBinUpEdge(ipt+1) - - rawy, effp, effnp, unc_rawy, unc_effp, unc_effnp = (np.zeros(n_sets) for _ in range(6)) - for iset, (hrawy, heffp, heffnp) in enumerate(zip(hist_rawy, hist_effp, hist_effnp)): - rawy.itemset(iset, hrawy.GetBinContent(ipt+1)) - effp.itemset(iset, heffp.GetBinContent(ipt+1)) - effnp.itemset(iset, heffnp.GetBinContent(ipt+1)) - unc_rawy.itemset(iset, hrawy.GetBinError(ipt+1)) - unc_effp.itemset(iset, heffp.GetBinError(ipt+1)) - unc_effnp.itemset(iset, heffnp.GetBinError(ipt+1)) + pt_min = hist_rawy[0].GetXaxis().GetBinLowEdge(ipt + 1) + pt_max = hist_rawy[0].GetXaxis().GetBinUpEdge(ipt + 1) + + rawy, effp, effnp, unc_rawy, unc_effp, unc_effnp = ( + np.zeros(n_sets) for _ in range(6) + ) + for iset, (hrawy, heffp, heffnp) in enumerate( + zip(hist_rawy, hist_effp, hist_effnp) + ): + rawy.itemset(iset, hrawy.GetBinContent(ipt + 1)) + effp.itemset(iset, heffp.GetBinContent(ipt + 1)) + effnp.itemset(iset, heffnp.GetBinContent(ipt + 1)) + unc_rawy.itemset(iset, hrawy.GetBinError(ipt + 1)) + unc_effp.itemset(iset, heffp.GetBinError(ipt + 1)) + unc_effnp.itemset(iset, heffnp.GetBinError(ipt + 1)) minimiser = CutVarMinimiser(rawy, effp, effnp, unc_rawy, unc_effp, unc_effnp) minimiser.minimise_system(cfg["minimisation"]["correlated"]) - hist_corry_prompt.SetBinContent(ipt+1, minimiser.get_prompt_yield_and_error()[0]) - hist_corry_prompt.SetBinError(ipt+1, minimiser.get_prompt_yield_and_error()[1]) - hist_corry_nonprompt.SetBinContent(ipt+1, minimiser.get_nonprompt_yield_and_error()[0]) - hist_corry_nonprompt.SetBinError(ipt+1, minimiser.get_nonprompt_yield_and_error()[1]) - hist_covariance.SetBinContent(ipt+1, minimiser.get_prompt_nonprompt_cov()) - hist_covariance.SetBinError(ipt+1, 0) - - canv_rawy, histos_rawy, leg_r = minimiser.plot_result(f"_pt{pt_min:.0f}_{pt_max:.0f}") + hist_corry_prompt.SetBinContent( + ipt + 1, minimiser.get_prompt_yield_and_error()[0] + ) + hist_corry_prompt.SetBinError( + ipt + 1, minimiser.get_prompt_yield_and_error()[1] + ) + hist_corry_nonprompt.SetBinContent( + ipt + 1, minimiser.get_nonprompt_yield_and_error()[0] + ) + hist_corry_nonprompt.SetBinError( + ipt + 1, minimiser.get_nonprompt_yield_and_error()[1] + ) + hist_covariance.SetBinContent(ipt + 1, minimiser.get_prompt_nonprompt_cov()) + hist_covariance.SetBinError(ipt + 1, 0) + + canv_rawy, histos_rawy, leg_r = minimiser.plot_result( + f"_pt{pt_min:.0f}_{pt_max:.0f}" + ) output.cd() canv_rawy.Write() for _, hist in histos_rawy.items(): hist.Write() - canv_eff, histos_eff, leg_e = minimiser.plot_efficiencies(f"_pt{pt_min:.0f}_{pt_max:.0f}") + canv_eff, histos_eff, leg_e = minimiser.plot_efficiencies( + f"_pt{pt_min:.0f}_{pt_max:.0f}" + ) output.cd() canv_eff.Write() for _, hist in histos_eff.items(): hist.Write() - canv_frac, histos_frac, leg_f = minimiser.plot_fractions(f"_pt{pt_min:.0f}_{pt_max:.0f}") + canv_frac, histos_frac, leg_f = minimiser.plot_fractions( + f"_pt{pt_min:.0f}_{pt_max:.0f}" + ) output.cd() canv_frac.Write() for _, hist in histos_frac.items(): @@ -117,21 +143,47 @@ def main(config): output_name_rawy_pdf = f"Distr_{cfg['output']['file'].replace('.root', '.pdf')}" output_name_eff_pdf = f"Eff_{cfg['output']['file'].replace('.root', '.pdf')}" output_name_frac_pdf = f"Frac_{cfg['output']['file'].replace('.root', '.pdf')}" - output_name_covmat_pdf = f"CovMatrix_{cfg['output']['file'].replace('.root', '.pdf')}" + output_name_covmat_pdf = ( + f"CovMatrix_{cfg['output']['file'].replace('.root', '.pdf')}" + ) if ipt == 0: - canv_rawy.SaveAs(f"{os.path.join(cfg['output']['directory'], output_name_rawy_pdf)}[") - canv_eff.SaveAs(f"{os.path.join(cfg['output']['directory'], output_name_eff_pdf)}[") - canv_frac.SaveAs(f"{os.path.join(cfg['output']['directory'], output_name_frac_pdf)}[") - canv_cov.SaveAs(f"{os.path.join(cfg['output']['directory'], output_name_covmat_pdf)}[") - canv_rawy.SaveAs(f"{os.path.join(cfg['output']['directory'], output_name_rawy_pdf)}") - canv_eff.SaveAs(f"{os.path.join(cfg['output']['directory'], output_name_eff_pdf)}") - canv_frac.SaveAs(f"{os.path.join(cfg['output']['directory'], output_name_frac_pdf)}") - canv_cov.SaveAs(f"{os.path.join(cfg['output']['directory'], output_name_covmat_pdf)}") - if ipt == hist_rawy[0].GetNbinsX()-1: - canv_rawy.SaveAs(f"{os.path.join(cfg['output']['directory'], output_name_rawy_pdf)}]") - canv_eff.SaveAs(f"{os.path.join(cfg['output']['directory'], output_name_eff_pdf)}]") - canv_frac.SaveAs(f"{os.path.join(cfg['output']['directory'], output_name_frac_pdf)}]") - canv_cov.SaveAs(f"{os.path.join(cfg['output']['directory'], output_name_covmat_pdf)}]") + canv_rawy.SaveAs( + f"{os.path.join(cfg['output']['directory'], output_name_rawy_pdf)}[" + ) + canv_eff.SaveAs( + f"{os.path.join(cfg['output']['directory'], output_name_eff_pdf)}[" + ) + canv_frac.SaveAs( + f"{os.path.join(cfg['output']['directory'], output_name_frac_pdf)}[" + ) + canv_cov.SaveAs( + f"{os.path.join(cfg['output']['directory'], output_name_covmat_pdf)}[" + ) + canv_rawy.SaveAs( + f"{os.path.join(cfg['output']['directory'], output_name_rawy_pdf)}" + ) + canv_eff.SaveAs( + f"{os.path.join(cfg['output']['directory'], output_name_eff_pdf)}" + ) + canv_frac.SaveAs( + f"{os.path.join(cfg['output']['directory'], output_name_frac_pdf)}" + ) + canv_cov.SaveAs( + f"{os.path.join(cfg['output']['directory'], output_name_covmat_pdf)}" + ) + if ipt == hist_rawy[0].GetNbinsX() - 1: + canv_rawy.SaveAs( + f"{os.path.join(cfg['output']['directory'], output_name_rawy_pdf)}]" + ) + canv_eff.SaveAs( + f"{os.path.join(cfg['output']['directory'], output_name_eff_pdf)}]" + ) + canv_frac.SaveAs( + f"{os.path.join(cfg['output']['directory'], output_name_frac_pdf)}]" + ) + canv_cov.SaveAs( + f"{os.path.join(cfg['output']['directory'], output_name_covmat_pdf)}]" + ) output.cd() hist_corry_prompt.Write() @@ -143,8 +195,12 @@ def main(config): if __name__ == "__main__": parser = argparse.ArgumentParser(description="Arguments") - parser.add_argument("config", metavar="text", default="config_cutvar_example.json", - help="JSON config file") + parser.add_argument( + "config", + metavar="text", + default="config_cutvar_example.json", + help="JSON config file", + ) args = parser.parse_args() main(args.config) diff --git a/PWGHF/D2H/Macros/cut_variation.py b/PWGHF/D2H/Macros/cut_variation.py index 2ab1aeb52c9..546e87403f2 100644 --- a/PWGHF/D2H/Macros/cut_variation.py +++ b/PWGHF/D2H/Macros/cut_variation.py @@ -8,6 +8,7 @@ """ import sys + import numpy as np import ROOT from style_formatter import set_global_style, set_object_style @@ -43,7 +44,7 @@ def __init__( # pylint: disable=too-many-arguments eff_nonprompt=None, unc_raw_yields=None, unc_eff_prompt=None, - unc_eff_nonprompt=None + unc_eff_nonprompt=None, ): self.raw_yields = raw_yields self.eff_prompt = eff_prompt @@ -68,7 +69,7 @@ def __init__( # pylint: disable=too-many-arguments self.m_corr_yields = None self.m_covariance = None - self.chi_2 = 0. + self.chi_2 = 0.0 self.ndf = self.n_sets - 2 def __check_input_consistency(self): @@ -79,7 +80,10 @@ def __check_input_consistency(self): self.n_sets = len(self.raw_yields) self.ndf = self.n_sets - 2 - if len(self.eff_prompt) != self.n_sets or len(self.eff_nonprompt) != self.n_sets: + if ( + len(self.eff_prompt) != self.n_sets + or len(self.eff_nonprompt) != self.n_sets + ): print("ERROR: number of raw yields and eff not consistent! Exit") sys.exit() @@ -87,7 +91,10 @@ def __check_input_consistency(self): print("ERROR: number of raw yield uncertainties not consistent! Exit") sys.exit() - if len(self.unc_eff_prompt) != self.n_sets or len(self.unc_eff_nonprompt) != self.n_sets: + if ( + len(self.unc_eff_prompt) != self.n_sets + or len(self.unc_eff_nonprompt) != self.n_sets + ): print("ERROR: number of raw yield uncertainties not consistent! Exit") sys.exit() @@ -118,7 +125,7 @@ def __initialise_objects(self): self.m_eff.itemset((i_set, 1), effnp) # pylint: disable=too-many-locals - def minimise_system(self, correlated=True, precision=1.e-8, max_iterations=100): + def minimise_system(self, correlated=True, precision=1.0e-8, max_iterations=100): """ Minimise the system of equations to compute the corrected yields @@ -144,11 +151,20 @@ def minimise_system(self, correlated=True, precision=1.e-8, max_iterations=100): zip(self.unc_raw_yields, self.unc_eff_prompt, self.unc_eff_nonprompt) ): for i_col, (rw_unc_col, effp_unc_col, effnp_unc_col) in enumerate( - zip(self.unc_raw_yields, self.unc_eff_prompt, self.unc_eff_nonprompt)): - unc_row = np.sqrt(rw_unc_row**2 + effp_unc_row**2 * self.m_corr_yields.item(0) - ** 2 + effnp_unc_row**2 * self.m_corr_yields.item(1)**2) - unc_col = np.sqrt(rw_unc_col**2 + effp_unc_col**2 * self.m_corr_yields.item(0) - ** 2 + effnp_unc_col**2 * self.m_corr_yields.item(1)**2) + zip( + self.unc_raw_yields, self.unc_eff_prompt, self.unc_eff_nonprompt + ) + ): + unc_row = np.sqrt( + rw_unc_row**2 + + effp_unc_row**2 * self.m_corr_yields.item(0) ** 2 + + effnp_unc_row**2 * self.m_corr_yields.item(1) ** 2 + ) + unc_col = np.sqrt( + rw_unc_col**2 + + effp_unc_col**2 * self.m_corr_yields.item(0) ** 2 + + effnp_unc_col**2 * self.m_corr_yields.item(1) ** 2 + ) if correlated and unc_row > 0 and unc_col > 0: if unc_row < unc_col: @@ -157,9 +173,9 @@ def minimise_system(self, correlated=True, precision=1.e-8, max_iterations=100): rho = unc_col / unc_row else: if i_row == i_col: - rho = 1. + rho = 1.0 else: - rho = 0. + rho = 0.0 cov_row_col = rho * unc_row * unc_col self.m_cov_sets.itemset((i_row, i_col), cov_row_col) @@ -172,13 +188,17 @@ def minimise_system(self, correlated=True, precision=1.e-8, max_iterations=100): self.m_covariance = np.linalg.inv(np.linalg.cholesky(self.m_covariance)) self.m_covariance = self.m_covariance.T * self.m_covariance - self.m_corr_yields = self.m_covariance * (m_eff_tr * self.m_weights) * self.m_rawy + self.m_corr_yields = ( + self.m_covariance * (m_eff_tr * self.m_weights) * self.m_rawy + ) self.m_res = self.m_eff * self.m_corr_yields - self.m_rawy m_res_tr = np.transpose(self.m_res) rel_delta = [ - (self.m_corr_yields.item(0)-m_corr_yields_old.item(0)) / self.m_corr_yields.item(0), - (self.m_corr_yields.item(1)-m_corr_yields_old.item(1)) / self.m_corr_yields.item(1) + (self.m_corr_yields.item(0) - m_corr_yields_old.item(0)) + / self.m_corr_yields.item(0), + (self.m_corr_yields.item(1) - m_corr_yields_old.item(1)) + / self.m_corr_yields.item(1), ] if rel_delta[0] < precision and rel_delta[1] < precision: @@ -193,17 +213,29 @@ def minimise_system(self, correlated=True, precision=1.e-8, max_iterations=100): for i_set, (effp, effnp) in enumerate(zip(self.eff_prompt, self.eff_nonprompt)): rawyp = effp * self.m_corr_yields.item(0) rawynp = effnp * self.m_corr_yields.item(1) - der_fp_p = (effp * (rawyp + rawynp) - effp**2 * self.m_corr_yields.item(0)) / (rawyp + rawynp)**2 - der_fp_np = -effp * effnp * self.m_corr_yields.item(0) / (rawyp + rawynp)**2 - der_fnp_np = (effnp * (rawyp + rawynp) - effnp**2 * self.m_corr_yields.item(1)) / (rawyp + rawynp)**2 - der_fnp_p = -effp * effnp * self.m_corr_yields.item(1) / (rawyp + rawynp)**2 - - unc_fp = np.sqrt(der_fp_p**2 * self.m_covariance.item(0, 0) + - der_fp_np**2 * self.m_covariance.item(1, 1) + - 2 * der_fp_p * der_fp_np * self.m_covariance.item(1, 0)) - unc_fnp = np.sqrt(der_fnp_p**2 * self.m_covariance.item(0, 0) + - der_fnp_np**2 * self.m_covariance.item(1, 1) + - 2 * der_fnp_p * der_fnp_np * self.m_covariance.item(1, 0)) + der_fp_p = ( + effp * (rawyp + rawynp) - effp**2 * self.m_corr_yields.item(0) + ) / (rawyp + rawynp) ** 2 + der_fp_np = ( + -effp * effnp * self.m_corr_yields.item(0) / (rawyp + rawynp) ** 2 + ) + der_fnp_np = ( + effnp * (rawyp + rawynp) - effnp**2 * self.m_corr_yields.item(1) + ) / (rawyp + rawynp) ** 2 + der_fnp_p = ( + -effp * effnp * self.m_corr_yields.item(1) / (rawyp + rawynp) ** 2 + ) + + unc_fp = np.sqrt( + der_fp_p**2 * self.m_covariance.item(0, 0) + + der_fp_np**2 * self.m_covariance.item(1, 1) + + 2 * der_fp_p * der_fp_np * self.m_covariance.item(1, 0) + ) + unc_fnp = np.sqrt( + der_fnp_p**2 * self.m_covariance.item(0, 0) + + der_fnp_np**2 * self.m_covariance.item(1, 1) + + 2 * der_fnp_p * der_fnp_np * self.m_covariance.item(1, 0) + ) self.frac_prompt.itemset(i_set, rawyp / (rawyp + rawynp)) self.frac_nonprompt.itemset(i_set, rawynp / (rawyp + rawynp)) self.unc_frac_prompt.itemset(i_set, unc_fp) @@ -265,32 +297,37 @@ def plot_result(self, suffix=""): ";cut set;raw yield", self.n_sets, -0.5, - self.n_sets-0.5 + self.n_sets - 0.5, ) hist_raw_yield_prompt = ROOT.TH1F( f"hRawYieldPromptVsCut{suffix}", ";cut set;prompt raw yield", self.n_sets, -0.5, - self.n_sets-0.5 + self.n_sets - 0.5, ) hist_raw_yield_nonprompt = ROOT.TH1F( f"hRawYieldNonPromptVsCut{suffix}", ";cut set;non-prompt raw yield", self.n_sets, -0.5, - self.n_sets-0.5 + self.n_sets - 0.5, ) hist_raw_yield_sum = ROOT.TH1F( f"hRawYieldSumVsCut{suffix}", ";cut set;raw yield", self.n_sets, -0.5, - self.n_sets-0.5 + self.n_sets - 0.5, ) for i_bin, (rawy, unc_rawy, effp, effnp) in enumerate( - zip(self.raw_yields, self.unc_raw_yields, self.eff_prompt, self.eff_nonprompt) + zip( + self.raw_yields, + self.unc_raw_yields, + self.eff_prompt, + self.eff_nonprompt, + ) ): hist_raw_yield.SetBinContent(i_bin + 1, rawy) hist_raw_yield.SetBinError(i_bin + 1, unc_rawy) @@ -299,8 +336,11 @@ def plot_result(self, suffix=""): unc_rawy_prompt = np.sqrt(self.m_covariance.item(0, 0)) * effp rawy_nonprompt = self.m_corr_yields.item(1) * effnp unc_rawy_nonprompt = np.sqrt(self.m_covariance.item(1, 1)) * effnp - unc_sum = np.sqrt(unc_rawy_prompt**2 + unc_rawy_nonprompt**2 + - 2 * self.m_covariance.item(1, 0) * effp * effnp) + unc_sum = np.sqrt( + unc_rawy_prompt**2 + + unc_rawy_nonprompt**2 + + 2 * self.m_covariance.item(1, 0) * effp * effnp + ) hist_raw_yield_prompt.SetBinContent(i_bin + 1, rawy_prompt) hist_raw_yield_prompt.SetBinError(i_bin + 1, unc_rawy_prompt) @@ -310,13 +350,20 @@ def plot_result(self, suffix=""): hist_raw_yield_sum.SetBinError(i_bin + 1, unc_sum) set_object_style(hist_raw_yield) - set_object_style(hist_raw_yield_prompt, color=ROOT.kRed+1, fillstyle=3145) - set_object_style(hist_raw_yield_nonprompt, color=ROOT.kAzure+4, fillstyle=3154) - set_object_style(hist_raw_yield_sum, color=ROOT.kGreen+2, fillstyle=0) + set_object_style(hist_raw_yield_prompt, color=ROOT.kRed + 1, fillstyle=3145) + set_object_style( + hist_raw_yield_nonprompt, color=ROOT.kAzure + 4, fillstyle=3154 + ) + set_object_style(hist_raw_yield_sum, color=ROOT.kGreen + 2, fillstyle=0) canvas = ROOT.TCanvas(f"cRawYieldVsCut{suffix}", "", 500, 500) - canvas.DrawFrame(-0.5, 0., self.n_sets - 0.5, hist_raw_yield.GetMaximum() * 1.2, - ";cut set;raw yield") + canvas.DrawFrame( + -0.5, + 0.0, + self.n_sets - 0.5, + hist_raw_yield.GetMaximum() * 1.2, + ";cut set;raw yield", + ) leg = ROOT.TLegend(0.6, 0.65, 0.8, 0.9) leg.SetBorderSize(0) leg.SetFillStyle(0) @@ -366,7 +413,7 @@ def plot_cov_matrix(self, correlated=True, suffix=""): padleftmargin=0.14, padbottommargin=0.12, padrightmargin=0.12, - palette=ROOT.kRainBow + palette=ROOT.kRainBow, ) hist_corr_matrix = ROOT.TH2F( @@ -374,10 +421,10 @@ def plot_cov_matrix(self, correlated=True, suffix=""): ";cut set;cut set", self.n_sets, -0.5, - self.n_sets-0.5, + self.n_sets - 0.5, self.n_sets, -0.5, - self.n_sets-0.5 + self.n_sets - 0.5, ) for i_row, unc_row in enumerate(self.unc_raw_yields): for i_col, unc_col in enumerate(self.unc_raw_yields): @@ -388,10 +435,10 @@ def plot_cov_matrix(self, correlated=True, suffix=""): rho = unc_col / unc_row else: if i_row == i_col: - rho = 1. + rho = 1.0 else: - rho = 0. - hist_corr_matrix.SetBinContent(i_row+1, i_col+1, rho) + rho = 0.0 + hist_corr_matrix.SetBinContent(i_row + 1, i_col + 1, rho) canvas = ROOT.TCanvas(f"cCorrMatrixCutSets{suffix}", "", 500, 500) hist_corr_matrix.Draw("colz") @@ -426,18 +473,23 @@ def plot_efficiencies(self, suffix=""): ";cut set;prompt acceptance #times efficiency", self.n_sets, -0.5, - self.n_sets-0.5 + self.n_sets - 0.5, ) hist_eff_nonprompt = ROOT.TH1F( f"hEffNonPromptVsCut{suffix}", ";cut set;non-prompt acceptance #times efficiency", self.n_sets, -0.5, - self.n_sets-0.5 + self.n_sets - 0.5, ) for i_bin, (effp, effnp, unc_effp, unc_effnp) in enumerate( - zip(self.eff_prompt, self.eff_nonprompt, self.unc_eff_prompt, self.unc_eff_nonprompt) + zip( + self.eff_prompt, + self.eff_nonprompt, + self.unc_eff_prompt, + self.unc_eff_nonprompt, + ) ): hist_eff_prompt.SetBinContent(i_bin + 1, effp) hist_eff_prompt.SetBinError(i_bin + 1, unc_effp) @@ -446,28 +498,30 @@ def plot_efficiencies(self, suffix=""): set_object_style( hist_eff_prompt, - color=ROOT.kRed+1, + color=ROOT.kRed + 1, fillstyle=0, markerstyle=ROOT.kFullCircle, - linewidth=2 + linewidth=2, ) set_object_style( hist_eff_nonprompt, - color=ROOT.kAzure+4, + color=ROOT.kAzure + 4, fillstyle=0, - markerstyle=ROOT.kFullSquare + markerstyle=ROOT.kFullSquare, ) canvas = ROOT.TCanvas(f"cEffVsCut{suffix}", "", 500, 500) - canvas.DrawFrame(-0.5, 1.e-5, self.n_sets - 0.5, 1., - ";cut set;acceptance #times efficiency") + canvas.DrawFrame( + -0.5, + 1.0e-5, + self.n_sets - 0.5, + 1.0, + ";cut set;acceptance #times efficiency", + ) canvas.SetLogy() hist_eff_prompt.Draw("esame") hist_eff_nonprompt.Draw("esame") - histos = { - "prompt": hist_eff_prompt, - "nonprompt": hist_eff_nonprompt - } + histos = {"prompt": hist_eff_prompt, "nonprompt": hist_eff_nonprompt} leg = ROOT.TLegend(0.2, 0.2, 0.4, 0.3) leg.SetBorderSize(0) leg.SetFillStyle(0) @@ -506,14 +560,14 @@ def plot_fractions(self, suffix=""): ";cut set;#it{f}_{prompt}", self.n_sets, -0.5, - self.n_sets-0.5 + self.n_sets - 0.5, ) hist_f_nonprompt = ROOT.TH1F( f"hFracNonPromptVsCut{suffix}", ";cut set;#it{f}_{non-prompt}", self.n_sets, -0.5, - self.n_sets-0.5 + self.n_sets - 0.5, ) for i_bin, (fracp, fracnp, unc_fp, unc_fnp) in enumerate( @@ -521,7 +575,7 @@ def plot_fractions(self, suffix=""): self.frac_prompt, self.frac_nonprompt, self.unc_frac_prompt, - self.unc_frac_nonprompt + self.unc_frac_nonprompt, ) ): hist_f_prompt.SetBinContent(i_bin + 1, fracp) @@ -531,27 +585,23 @@ def plot_fractions(self, suffix=""): set_object_style( hist_f_prompt, - color=ROOT.kRed+1, + color=ROOT.kRed + 1, fillstyle=0, markerstyle=ROOT.kFullCircle, - linewidth=2 + linewidth=2, ) set_object_style( hist_f_nonprompt, - color=ROOT.kAzure+4, + color=ROOT.kAzure + 4, fillstyle=0, - markerstyle=ROOT.kFullSquare + markerstyle=ROOT.kFullSquare, ) canvas = ROOT.TCanvas(f"cFracVsCut{suffix}", "", 500, 500) - canvas.DrawFrame(-0.5, 0., self.n_sets - 0.5, 1., - ";cut set;fraction") + canvas.DrawFrame(-0.5, 0.0, self.n_sets - 0.5, 1.0, ";cut set;fraction") hist_f_prompt.Draw("esame") hist_f_nonprompt.Draw("esame") - histos = { - "prompt": hist_f_prompt, - "nonprompt": hist_f_nonprompt - } + histos = {"prompt": hist_f_prompt, "nonprompt": hist_f_nonprompt} leg = ROOT.TLegend(0.2, 0.8, 0.4, 0.9) leg.SetBorderSize(0) leg.SetFillStyle(0) diff --git a/PWGHF/D2H/Macros/style_formatter.py b/PWGHF/D2H/Macros/style_formatter.py index bb3d98d9640..02cc52507c4 100644 --- a/PWGHF/D2H/Macros/style_formatter.py +++ b/PWGHF/D2H/Macros/style_formatter.py @@ -261,7 +261,7 @@ def set_object_style(obj, **kwargs): if "markersize" in kwargs: obj.SetMarkerSize(kwargs["markersize"]) else: - obj.SetMarkerSize(1.) + obj.SetMarkerSize(1.0) if "markerstyle" in kwargs: obj.SetMarkerStyle(kwargs["markerstyle"])