[WIP] Adding options for restrict cards

mcgeneration/configure_gridpack.py

@@ -457,7 +457,8 @@ def main():
     random.seed()
     stype = ScanType.FROMFILE
     btype = BatchType.CMSCONNECT
-    tag   = 'Run3_52WCs_SMEFTsim_top'
+    tag   = 'Example1'
+    restrict = False 
     runs  = 1               # if set to 0, will only make a single gridpack
     npts  = 0
     #scan_files = [
@@ -565,7 +566,7 @@ def main():
     sm_pt    = {}
     for dof in dof_list: sm_pt[dof.getName()] = 0.0
 
-    gridpack = Gridpack(stype=stype,btype=btype,default_limits=[-20.0,20.0])
+    gridpack = Gridpack(stype=stype,btype=btype,default_limits=[-20.0,20.0],restrict=restrict)
     gridpack.setOptions(runcard_ops=rc_ops)
     # For using a different model
     gridpack.setOptions(coupling_string="SMHLOOP=0 NP=1 NPprop=0",replace_model=["SMEFTsim_topU3l_MwScheme_UFO","SMEFTsim_top_MwScheme_UFO"])

mcgeneration/helpers/Gridpack.py

@@ -65,6 +65,7 @@ def __init__(self,**kwargs):
             'replace_model': None,          # If not None overwrites the import model line of the process card
             'flavor_scheme': 5,
             'default_limits': [-10,10],
+            'restrict': False,
         }
 
         self.setOptions(**kwargs)
@@ -203,6 +204,9 @@ def saveProcessCard(self,indent=0):
         if self.ops['replace_model']:
             old = self.ops['replace_model'][0]
             new = self.ops['replace_model'][1]
+            if self.ops['restrict']: new += f'-massless_{self.ops["process"]}'
+            if len(list(self.ops['coeffs'])) == 1:
+                new += f'_{list(self.ops["coeffs"])[0]}'
             print("{ind}Using {model} model".format(model=new,ind=indent_str))
             sed_str = "s|import model {old}|import model {new}|g".format(old=old,new=new)
             subprocess.Popen(['sed','-i','-e',sed_str,fpath]).communicate()
@@ -229,6 +233,12 @@ def saveReweightCard(self):
 
         save_scan_points(scanfile,self.ops['coeffs'],self.scan_pts)
         make_reweight_card(rwgt_tar,self.ops['coeffs'],self.scan_pts)
+        if self.ops['restrict']:
+            mpath  =  'addons/models/SMEFTsim_top_MwScheme_UFO/restrict_massless.dat'
+            ompath = f'addons/models/SMEFTsim_top_MwScheme_UFO/restrict_massless_{self.ops["process"]}.dat'
+            if len(list(self.ops['coeffs'])) == 1:
+                ompath = ompath[:-4] + f'_{list(self.ops["coeffs"])[0]}.dat'
+            make_restrict_card(mpath, ompath, keep=True, SMEFT=list(self.ops['coeffs']))
 
         return rwgt_tar
 

mcgeneration/helpers/helper_tools.py

@@ -96,6 +96,73 @@ def make_reweight_card(file_name,dofs,pts):
                     f.write("\nset %s %.6f" % (k2,v2))
             f.write("\n")
 
+def make_restrict_card(file_name,out_fname,keep=True,**blocks):
+    '''
+        file_name: Name of the restrict card to read, which will serve as the basis for the new restrict card
+        out_fname: Name of the modified restrict card
+        keep: If true, then for a given block, the listed parameters will be NOT zero'd out.
+              If false, then for a given block, the listed parameters will be zero'd out.
+        blocks: {
+            "block_A" : ["param1"."param2", ...],
+            "block_B" : [ ... ],
+        }
+    '''
+    counter = 1
+    indent = " "*2
+    lines = []
+    block = None
+    with open(file_name,'r') as f:
+        for l in f.readlines():
+            # Check if this ENTIRE line is a comment
+            is_comment = l.startswith("#")
+            # Check if this line specifies the start of a new LHA block section
+            is_block_header = l.lower().startswith("block")
+            # Check if this line is an empty line
+            is_empty_line = len(l) == 0
+            # Skip lines we know we won't need to edit
+            skip = is_comment or is_block_header or is_empty_line
+
+            if is_block_header:
+                # Store the name of the current block
+                block = l.split()[1]
+
+            if l.lower().startswith("decay"):
+                # Decay lines are their own thing separate from LHA block stuff, so don't mess with them
+                skip = True
+            elif block == "QNUMBERS":
+                # QNUMBERS blocks have a bit different syntax then other blocks, so avoid them as well
+                skip = True
+
+
+            # Avoid dealing with lines that should never need to be edited
+            if not skip:
+                x = l.split(' # ')
+                if len(x) == 1:
+                    continue
+                data, param_name = [x.strip() for x in l.split(" # ")]
+                # data should always be 2 numbers separated by a single space
+                idx, value = data.split()
+                if block in blocks:
+                    params = blocks[block]
+                    if keep:
+                        # Zero out any params that aren't specified
+                        if param_name in params:
+                            l = f"{indent}{idx:>3} 0.{counter:0>7}e+00 # {param_name}"
+                            counter += 1
+                        else:
+                            l = f"{indent}{idx:>3} 0.0000000e+00 # {param_name}"
+                    else:
+                        # Zero out any params that are specified
+                        if param_name in params:
+                            l = f"{indent}{idx:>3} 0.0000000e+00 # {param_name}"
+                        else:
+                            l = f"{indent}{idx:>3} 0.{counter:0>7}e+00 # {param_name}"
+                            counter += 1
+            # The new restrict card should be (as far as lines go) a 1-to-1 mirror of the base card
+            lines.append(l.rstrip())
+    with open(out_fname,'w') as f:
+        f.write("\n".join(lines))
+
 # Reads a limit file and returns a dictionary mapping the WCs to their respective high,low limits to use
 def parse_limit_file(fpath):
     wc_limits = {}