[PWGHF] Add methods to compute fraction, provided both efficiencies and cut-variation results externally (AliceO2Group#9616)

Author: fgrosa

PWGHF/D2H/Macros/compute_fraction_cutvar.py

@@ -10,11 +10,13 @@
 import argparse
 import json
 import os
+import sys
 
 import numpy as np  # pylint: disable=import-error
 import ROOT  # pylint: disable=import-error
+sys.path.insert(0, '..')
 from cut_variation import CutVarMinimiser
-from style_formatter import set_object_style
+from utils.style_formatter import set_object_style
 
 # pylint: disable=no-member,too-many-locals,too-many-statements
 
@@ -176,24 +178,39 @@ def main(config):
         canv_cov.Write()
         histo_cov.Write()
 
+        canv_combined = ROOT.TCanvas("canv_combined", "", 1000, 1000)
+        canv_combined.Divide(2, 2)
+        canv_combined.cd(1)
+        canv_rawy.DrawClonePad()
+        canv_combined.cd(2)
+        canv_eff.DrawClonePad()
+        canv_combined.cd(3)
+        canv_frac.DrawClonePad()
+        canv_combined.cd(4)
+        canv_cov.DrawClonePad()
+
         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_pdf = f"{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_combined.SaveAs(f"{os.path.join(cfg['output']['directory'], output_name_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_combined.SaveAs(f"{os.path.join(cfg['output']['directory'], output_name_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_combined.SaveAs(f"{os.path.join(cfg['output']['directory'], output_name_pdf)}]")
 
     output.cd()
     hist_corry_prompt.Write()

PWGHF/D2H/Macros/cut_variation.py

@@ -11,7 +11,8 @@
 
 import numpy as np  # pylint: disable=import-error
 import ROOT  # pylint: disable=import-error
-from style_formatter import set_global_style, set_object_style
+sys.path.insert(0, '..')
+from utils.style_formatter import set_global_style, set_object_style
 
 
 # pylint: disable=too-many-instance-attributes
@@ -295,6 +296,49 @@ def get_raw_prompt_fraction(self, effacc_p, effacc_np):
 
         return f_p, f_p_unc
 
+    def get_raw_prompt_fraction_ext(self, corry_p, corry_np, unc_corry_p,
+                                    unc_corry_np, cov_p_np, effacc_p, effacc_np):
+        """
+        Helper function to get the raw prompt fraction given the efficiencies
+
+        Parameters
+        -----------------------------------------------------
+        - corry_p: float
+            corrected yield for prompt signal
+        - corry_np: float
+            corrected yield for non-prompt signal
+        - unc_corry_np: float
+            uncertainty on corrected yield for prompt signal
+        - unc_corry_np: float
+            uncertainty on corrected yield for non-prompt signal
+        - cov_p_np: float
+            covariance between prompt and non-prompt signal
+        - effacc_p: float
+            eff x acc for prompt signal
+        - effacc_np: float
+            eff x acc for non-prompt signal
+
+        Returns
+        -----------------------------------------------------
+        - f_p, f_p_unc: (float, float)
+            raw prompt fraction with its uncertainty
+        """
+
+        rawy_p = effacc_p * corry_p
+        rawy_np = effacc_np * corry_np
+        f_p = rawy_p / (rawy_p + rawy_np)
+
+        # derivatives of prompt fraction wrt corr yields
+        d_p = (effacc_p * (rawy_p + rawy_np) - effacc_p**2 * corry_p) / (rawy_p + rawy_np) ** 2
+        d_np = -effacc_np * rawy_p / (rawy_p + rawy_np) ** 2
+        f_p_unc = np.sqrt(
+            d_p**2 * unc_corry_p**2
+            + d_np**2 * unc_corry_np**2
+            + 2 * d_p * d_np * cov_p_np
+        )
+
+        return f_p, f_p_unc
+
     def get_raw_nonprompt_fraction(self, effacc_p, effacc_np):
         """
         Helper function to get the raw non-prompt fraction given the efficiencies
@@ -318,6 +362,41 @@ def get_raw_nonprompt_fraction(self, effacc_p, effacc_np):
 
         return f_np, f_np_unc
 
+    def get_raw_nonprompt_fraction_ext(self, corry_p, corry_np, unc_corry_p,
+                                       unc_corry_np, cov_p_np, effacc_p, effacc_np):
+        """
+        Helper function to get the raw non-prompt fraction given the efficiencies
+
+        Parameters
+        -----------------------------------------------------
+        - corry_p: float
+            corrected yield for prompt signal
+        - corry_np: float
+            corrected yield for non-prompt signal
+        - unc_corry_np: float
+            uncertainty on corrected yield for prompt signal
+        - unc_corry_np: float
+            uncertainty on corrected yield for non-prompt signal
+        - cov_p_np: float
+            covariance between prompt and non-prompt signal
+        - effacc_p: float
+            eff x acc for prompt signal
+        - effacc_np: float
+            eff x acc for non-prompt signal
+
+        Returns
+        -----------------------------------------------------
+        - f_np, f_np_unc: (float, float)
+            raw non-prompt fraction with its uncertainty
+
+        """
+
+        f_p, f_np_unc = self.get_raw_prompt_fraction_ext(corry_p, corry_np, unc_corry_p,
+                                                         unc_corry_np, cov_p_np, effacc_p, effacc_np)
+        f_np = 1 - f_p
+
+        return f_np, f_np_unc
+
     def get_corr_prompt_fraction(self):
         """
         Helper function to get the corrected prompt fraction
@@ -365,7 +444,7 @@ def plot_result(self, suffix=""):
             needed otherwise it is destroyed
         """
 
-        set_global_style(padleftmargin=0.16, padbottommargin=0.12, titleoffsety=1.6)
+        set_global_style(padleftmargin=0.16, padbottommargin=0.12, padtopmargin=0.075, titleoffsety=1.6)
 
         hist_raw_yield = ROOT.TH1F(
             f"hRawYieldVsCut{suffix}",
@@ -435,7 +514,7 @@ def plot_result(self, suffix=""):
             hist_raw_yield.GetMaximum() * 1.2,
             ";cut set;raw yield",
         )
-        leg = ROOT.TLegend(0.6, 0.65, 0.8, 0.9)
+        leg = ROOT.TLegend(0.6, 0.65, 0.8, 0.85)
         leg.SetBorderSize(0)
         leg.SetFillStyle(0)
         leg.SetTextSize(0.04)