Add the possibility of setting fixed sigma

Author: saganatt

machine_learning_hep/analysis/analyzerdhadrons.py

@@ -100,6 +100,8 @@ def __init__(self, datap, case, typean, period):
         self.n_fileff = os.path.join(self.d_resultsallpmc, self.n_fileff)
         self.p_bin_width = datap["analysis"][self.typean]["bin_width"]
         self.p_rebin = datap["analysis"][self.typean]["n_rebin"]
+        self.p_fixed_sigma = datap["analysis"][self.typean]["fixed_sigma"]
+        self.p_fixed_sigma_val = datap["analysis"][self.typean]["fixed_sigma_val"]
         self.p_pdfnames = datap["analysis"][self.typean]["pdf_names"]
         self.p_param_names = datap["analysis"][self.typean]["param_names"]
 
@@ -169,11 +171,14 @@ def _save_hist(self, hist, filename, option=""):
         self.rfigfile.WriteObject(hist, rfilename)
 
     # region fitting
-    def _roofit_mass(self, level, hist, ipt, pdfnames, param_names, fitcfg, roows=None, filename=None):
+    def _roofit_mass(self, level, hist, ipt, pdfnames, param_names, fitcfg, fixed_sigma, fixed_sigma_val, # pylint: disable=too-many-arguments
+                     roows=None, filename=None):
         if fitcfg is None:
             return None, None, None, None, None
         try:
-            res, ws, frame, residual_frame, data_hist, model = self.fitter.fit_mass_new(hist, pdfnames, fitcfg, level,
+            res, ws, frame, residual_frame, data_hist, model = self.fitter.fit_mass_new(hist, pdfnames, param_names,
+                                                                                        fitcfg, level,
+                                                                                        fixed_sigma, fixed_sigma_val,
                                                                                         roows=roows, plot=True)
         except ValueError:
             self.logger.error(f"Could not do fitting on {level} for pt {self.bins_candpt[ipt]} - {self.bins_candpt[ipt+1]}")
@@ -356,6 +361,8 @@ def fit(self):
                             self.p_pdfnames,
                             self.p_param_names,
                             fitcfg,
+                            self.p_fixed_sigma[ipt],
+                            self.p_fixed_sigma_val[ipt],
                             roows,
                             f"roofit/h_mass_fitted_pthf-{ptrange[0]}-{ptrange[1]}_{level}.png",
                         )

machine_learning_hep/analysis/analyzerdhadrons_mult.py

@@ -105,6 +105,8 @@ def __init__(self, datap, case, typean, period):
         self.p_bin_width = datap["analysis"][self.typean]["bin_width"]
 
         self.p_rebin = datap["analysis"][self.typean]["n_rebin"]
+        self.p_fixed_sigma = datap["analysis"][self.typean]["fixed_sigma"]
+        self.p_fixed_sigma_val = datap["analysis"][self.typean]["fixed_sigma_val"]
         self.p_pdfnames = datap["analysis"][self.typean]["pdf_names"]
         self.p_param_names = datap["analysis"][self.typean]["param_names"]
 
@@ -197,11 +199,14 @@ def _save_hist(self, hist, filename, option=""):
         self.rfigfile.WriteObject(hist, rfilename)
 
     # region fitting
-    def _roofit_mass(self, level, hist, ipt, pdfnames, param_names, fitcfg, roows=None, filename=None):
+    def _roofit_mass(self, level, hist, ipt, pdfnames, param_names, fitcfg, fixed_sigma, fixed_sigma_val, # pylint: disable=too-many-arguments
+                     roows=None, filename=None):
         if fitcfg is None:
             return None, None, None, None, None
         try:
-            res, ws, frame, residual_frame, data_hist, model = self.fitter.fit_mass_new(hist, pdfnames, fitcfg, level,
+            res, ws, frame, residual_frame, data_hist, model = self.fitter.fit_mass_new(hist, pdfnames, param_names,
+                                                                                        fitcfg, level,
+                                                                                        fixed_sigma, fixed_sigma_val,
                                                                                         roows=roows, plot=True)
         except ValueError:
             self.logger.error(f"Could not do fitting on {level} for pt {self.bins_candpt[ipt]} - {self.bins_candpt[ipt+1]}")
@@ -406,6 +411,8 @@ def fit(self):
                                 self.p_pdfnames,
                                 self.p_param_names,
                                 fitcfg,
+                                self.p_fixed_sigma[ipt],
+                                self.p_fixed_sigma_val[ipt],
                                 roows,
                                 f"roofit/mult_{multrange[0]}-{multrange[1]}/"
                                 f"h_mass_fitted_pthf-{ptrange[0]}-{ptrange[1]}"

machine_learning_hep/fitting/roofitter.py

@@ -34,7 +34,8 @@ def __init__(self):
 
     # pylint: disable=too-many-branches
     def fit_mass_new(
-        self, hist, pdfnames: dict, fit_spec: dict, level: str, roows: ROOT.RooWorkspace = None, plot: bool = False
+        self, hist, pdfnames: dict, param_names: dict, fit_spec: dict, level: str,
+        fixed_sigma: bool = False, fixed_sigma_val: float = 0., roows: ROOT.RooWorkspace = None, plot: bool = False
     ):
         """New fit method"""
         if hist.GetEntries() == 0:
@@ -54,6 +55,11 @@ def fit_mass_new(
             raise ValueError("model not set")
 
         m = ws.var(var_m)
+        if level == "mc":
+            sigma_sgn = ws.var(param_names["gauss_sigma"])
+            if fixed_sigma:
+                sigma_sgn.setVal(fixed_sigma_val)
+                sigma_sgn.setConstant(True)
 
         if level == "data" and USE_EXTMODEL:
             signal_pdf = ws.pdf(pdfnames["pdf_sig"])