Merge pull request #996 from cms-analysis/impacts_upd_regex

Update Impacts.py
cms-analysis · Jul 24, 2024 · 0cb4611 · 0cb4611
2 parents e5aa5f7 + c534f67
commit 0cb4611
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 27 deletions.
diff --git a/python/tool_base/FastScan.py b/python/tool_base/FastScan.py
@@ -42,13 +42,6 @@ def attach_args(self, group):
         group.add_argument("-o", "--output", default="nll", help="Name of the output file, without the .pdf extension")
         group.add_argument("-p", "--points", default=200, type=int, help="Number of NLL points to sample in each scan")
 
-    def RooColIter(self, coll):
-        it = coll.createIterator()
-        var = it.Next()
-        while var:
-            yield var
-            var = it.Next()
-
     def run_method(self):
         ROOT.gROOT.SetBatch(ROOT.kTRUE)
         outfile = ROOT.TFile("%s.root" % self.args.output, "RECREATE")
@@ -83,7 +76,7 @@ def run_method(self):
         page = 0
         doPars = []
 
-        for par in self.RooColIter(pars):
+        for par in pars:
             if par.isConstant():
                 continue
             if self.args.match is not None:

diff --git a/python/tool_base/Impacts.py b/python/tool_base/Impacts.py
@@ -33,7 +33,6 @@ def attach_intercept_args(self, group):
             --output stages. Note the ordering of POIs in the list must also be
             identical in each step.""",
         )
-        group.add_argument("--setPhysicsModelParameters")
         group.add_argument("--setParameters")
         group.add_argument("--name", "-n", default="Test")
 
@@ -101,11 +100,8 @@ def run_method(self):
         # Put intercepted args back
         passthru.extend(["-m", mh])
         passthru.extend(["-d", ws])
-        if self.args.setPhysicsModelParameters is not None:
-            passthru.extend(["--setPhysicsModelParameters", self.args.setPhysicsModelParameters])
         if self.args.setParameters is not None:
             passthru.extend(["--setParameters", self.args.setParameters])
-            self.args.setPhysicsModelParameters = self.args.setParameters
         pass_str = " ".join(passthru)
 
         paramList = []
@@ -204,7 +200,26 @@ def run_method(self):
 
         print("Have parameters: " + str(len(paramList)))
 
-        prefit = utils.prefit_from_workspace(ws, "w", paramList, self.args.setPhysicsModelParameters)
+        varList = utils.list_from_workspace(ws, "w", "variables")
+        if self.args.setParameters is not None:
+            set_parameters = self.args.setParameters.split(",")
+            set_parameters_str = ""
+            for ind, setParam in enumerate(set_parameters):
+                if "rgx{" in setParam:
+                    eqs_to = setParam.split("=")[-1]
+                    pattern = setParam.split("=")[0]
+                    pattern = pattern.replace("'rgx{", "").replace("}'", "")
+                    pattern = pattern.replace("rgx{", "").replace("}", "")
+                    set_parameters[ind] = ""
+                    for var in varList:
+                        if re.search(pattern, var):
+                            var_str = var + "=" + eqs_to
+                            set_parameters_str += var_str + ","
+                else:
+                    set_parameters_str += setParam + ","
+            self.args.setParameters = set_parameters_str.rstrip(",")
+
+        prefit = utils.prefit_from_workspace(ws, "w", paramList, self.args.setParameters)
         res = {}
         if not self.args.noInitialFit:
             res["POIs"] = []
@@ -264,15 +279,9 @@ def run_method(self):
             print("Missing inputs: " + ",".join(missing))
 
     def all_free_parameters(self, file, wsp, mc, pois):
-        res = []
         wsFile = ROOT.TFile.Open(file)
         w = wsFile.Get(wsp)
         config = w.genobj(mc)
         pdfvars = config.GetPdf().getParameters(config.GetObservables())
-        it = pdfvars.createIterator()
-        var = it.Next()
-        while var:
-            if var.GetName() not in pois and (not var.isConstant()) and var.InheritsFrom("RooRealVar"):
-                res.append(var.GetName())
-            var = it.Next()
+        res = [var.GetName() for var in pdfvars if (var.GetName() not in pois and (not var.isConstant()) and var.InheritsFrom("RooRealVar"))]
         return res
diff --git a/python/tool_base/utils.py b/python/tool_base/utils.py
@@ -33,15 +33,10 @@ def split_vals(vals, fmt_spec=None):
 
 def list_from_workspace(file, workspace, set):
     """Create a list of strings from a RooWorkspace set"""
-    res = []
     wsFile = ROOT.TFile(file)
     ws = wsFile.Get(workspace)
-    argSet = ws.set(set)
-    it = argSet.createIterator()
-    var = it.Next()
-    while var:
-        res.append(var.GetName())
-        var = it.Next()
+    argSet = ws.allVars() if set == "variables" else ws.set(set)
+    res = [var.GetName() for var in argSet]
     return res