Update formula when modifying workspace

If in workspace there is an object defined as RooFormulaVar::SigXsecOverSM_ggZHbb_kappa[ actualVars=(kappaL,kappa2V,kappat,den_mu) formula="(1.+(x[0]-1.0)*0.0)/x[3]" ] = 0.999862 (bbtautau kappa_lambda workspace for example), we are not able to simply rename variables used in these formula: if we rename 'kappaL' to 'klambda', the name of variables in formula won't change, so we will get error saying that

RooFormula 'SigXsecOverSM_ggZHbb_kappa' did not compile or is invalid.
Input:
        (1.+(kappaL-1.0)*0.0)/kappat
Passed over to TFormula:
        (1.+(kappaL-1.0)*0.0)/x[3]

This is something similar as mentioned here https://gitlab.cern.ch/atlas-physics/HDBS/DiHiggs/combination/hh_combination_fw/-/blob/master/deprecated/doc/troubleshooting.md?ref_type=heads#symptoms, however, since whether to specify the variables in RooFormulaVar or not is just a choice, I don't think the requirement on inputs make a lot of sense, instead, we should be able to update RooFormulaVar on the fly.

quickstats/components/workspaces/xml_ws_modifier.py

@@ -270,6 +270,48 @@ class XMLWSModifier(XMLWSBase):
         if len(self.actions['rename']['model_config']) > 1:
             raise RuntimeError("model config is renamed more than once")
     
+    def correcct_formula(self,ws_orig, ws_tmp, rename_map):
+        old_var_expr = ",".join(list(rename_map.keys()))
+        new_var_expr = ",".join(list(rename_map.values()))
+        tested_vars = []
+        iter = ws_orig.componentIterator()
+        obj = iter.Next()
+        while obj:
+            if isinstance(obj, ROOT.RooFormulaVar):
+                formula = obj.expression()
+                formula_bak = formula
+                actual_vars = obj.dependents()
+                for ivar,var in enumerate(actual_vars):
+                    formula = formula.replace(var.GetName(), f'@{ivar}')
+                    formula = formula.replace(f'x[{ivar}]', f'@{ivar}')
+
+                # Sometimes the input variables will be a function of other variables
+                for actual_var in actual_vars:
+                    if actual_var in tested_vars: continue
+                    tested_vars.append(actual_var)
+                    if isinstance(actual_var, ROOT.RooAddition):
+                        component_list = actual_var.getComponents()
+                        add_iter = component_list.createIterator()
+                        add_component = add_iter.Next()
+                        while add_component:
+                            if isinstance(add_component, ROOT.RooFormulaVar):
+                                add_formula = add_component.expression()
+                                add_formula_bak = add_formula
+                                add_actual_vars = add_component.dependents()
+                                for add_ivar,add_var in enumerate(add_actual_vars):
+                                    add_formula = add_formula.replace(add_var.GetName(), f'@{add_ivar}')
+                                    add_formula = add_formula.replace(f'x[{add_ivar}]', f'@{add_ivar}')
+                                if add_formula!=add_formula_bak:
+                                    add_new_formula_var = ROOT.RooFormulaVar(add_component.GetName(), add_component.GetTitle(), add_formula, add_actual_vars)
+                                    getattr(ws_tmp, "import")(add_new_formula_var, ROOT.RooFit.RenameVariable(old_var_expr, new_var_expr), ROOT.RooFit.RecycleConflictNodes())
+                            add_component = add_iter.Next()
+                
+                # Update the formula only if there is any change
+                if formula!=formula_bak:
+                    new_formula_var = ROOT.RooFormulaVar(obj.GetName(), obj.GetTitle(), formula, actual_vars)
+                    getattr(ws_tmp, "import")(new_formula_var, ROOT.RooFit.RenameVariable(old_var_expr, new_var_expr), ROOT.RooFit.RecycleConflictNodes())
+            obj = iter.Next()
+    
     def create_modified_workspace(self, infile:Optional[str]=None, outfile:Optional[str]=None,
                                   import_class_code:bool=True, recreate:bool=True):
         with Timer() as t:
@@ -316,6 +358,10 @@ class XMLWSModifier(XMLWSBase):
             self.stdout.info(title_str, bare=True)
             rename_map = self.get_rename_map(ws_orig, ws_tmp)
             actual_rename_map = self.get_actual_rename_map(rename_map, ws_orig, ws_tmp)
+
+            # In case the variable exist in RooFormulaVar
+            self.correcct_formula(ws_orig, ws_tmp, rename_map)
+
             self.rename_variables(mc_orig, ws_tmp, rename_map)
             if not rename_map:
                 self.stdout.info("No variables are renamed.")