[WIP] Adding options for restrict cards

README.md

@@ -31,6 +31,24 @@ gp = Gridpack(replace_model=['dim6top_LO_UFO','smloop'])
 # Alternatively via setOptions
 gp.setOptions(replace_model=['dim6top_LO_UFO','smloop'])
 ```
+- The `restrict` option will tell the code to create a restrict card for the model you choose to use. The value you pass when setting the `restrict` option is a dictionary with the following structure. At the moment, the `restrict` option only has an effect when using `ScanType.SLINSPACE`.
+```python
+restrict = {
+    "ref": "restrict_massless.dat",
+    "blocks": ["SMEFT","SMEFTcpv"],
+    "keep": True
+}
+gp.setOptions(restrict=restrict)
+```
+`ref`
+: This should be the name of a restrict card that is already located in the directory of the model you intend to use. It will be used as a template reference to create the actual restrict card that Madgraph will be pointed to.
+
+`blocks`
+: This is a list of UFO block names that the code will consider when deciding to keep or exclude certain parameters. Any parameters specified under other blocks will be kept at the values they had in the reference restrict card. The block names _are_ case sensitive!
+
+`keep`
+: A booleon value that indicates how to treat matched parameters. If `True`, then only matched parameters are kept and all others in the block are set to 0. If `False`, then only the matched parameters are set to 0.
+
 
 # Resubmitting failed gridpacks:
 Sometimes a gridpack will get past the `codegen` stage but fail in the `integrate`, resulting in a `.tar.xz` file less than 10 MB. <br>

mcgeneration/addons/limits/SMEFTsim_top_limits.txt

@@ -0,0 +1,141 @@
+tllq4fNoSchanWNoHiggs0p_cQe1     -92.96  93.19  
+tllq4fNoSchanWNoHiggs0p_cQe2     -92.96  93.19  
+tllq4fNoSchanWNoHiggs0p_cQe3     -92.96  93.19  
+tllq4fNoSchanWNoHiggs0p_cQl31    -28.4   27.92  
+tllq4fNoSchanWNoHiggs0p_cQl32    -28.4   27.92  
+tllq4fNoSchanWNoHiggs0p_cQl33    -28.4   27.92  
+tllq4fNoSchanWNoHiggs0p_cQl11    -94.39  91.73  
+tllq4fNoSchanWNoHiggs0p_cQl12    -94.39  91.73  
+tllq4fNoSchanWNoHiggs0p_cQl13    -94.39  91.73  
+tllq4fNoSchanWNoHiggs0p_cbWRe      -6.67   6.67   
+tllq4fNoSchanWNoHiggs0p_cHQ3     -13.93  8.09   
+tllq4fNoSchanWNoHiggs0p_cHQ1     -70.28  50.89  
+tllq4fNoSchanWNoHiggs0p_cHtb      -96.82  83.65  
+tllq4fNoSchanWNoHiggs0p_cHtbRe     -25.33  25.33  
+tllq4fNoSchanWNoHiggs0p_ctG      -157.91 157.91 
+tllq4fNoSchanWNoHiggs0p_ctWRe      -5.88   5.53   
+tllq4fNoSchanWNoHiggs0p_ctBRe      -13.36  13.71  
+tllq4fNoSchanWNoHiggs0p_cte1     -157.91 157.91 
+tllq4fNoSchanWNoHiggs0p_cte2     -157.91 157.91 
+tllq4fNoSchanWNoHiggs0p_cte3     -157.91 157.91 
+tllq4fNoSchanWNoHiggs0p_cleQt1Re1    -157.91 157.91 
+tllq4fNoSchanWNoHiggs0p_cleQt1Re2    -157.91 157.91 
+tllq4fNoSchanWNoHiggs0p_cleQt1Re3    -157.91 157.91 
+tllq4fNoSchanWNoHiggs0p_cleQt3Re1    -22.84  22.84  
+tllq4fNoSchanWNoHiggs0p_cleQt3Re2    -22.84  22.84  
+tllq4fNoSchanWNoHiggs0p_cleQt3Re3    -22.84  22.84  
+tllq4fNoSchanWNoHiggs0p_ctl     -157.91 157.91 
+tllq4fNoSchanWNoHiggs0p_ctHRe      -78.02  80.42  
+ttHJet_cQe1      -157.91 157.91 
+ttHJet_cQe2      -157.91 157.91 
+ttHJet_cQe3      -157.91 157.91 
+ttHJet_cQl31     -157.91 157.91 
+ttHJet_cQl32     -157.91 157.91 
+ttHJet_cQl33     -157.91 157.91 
+ttHJet_cQl11     -157.91 157.91 
+ttHJet_cQl12     -157.91 157.91 
+ttHJet_cQl13     -157.91 157.91 
+ttHJet_cbWRe       -46.59  46.59  
+ttHJet_cHQ3      -41.29  40.46  
+ttHJet_cHQ1      -157.91 157.91 
+ttHJet_cHtb       -157.91 157.91 
+ttHJet_cHtbRe      -108.22 108.22 
+ttHJet_ctG       -2.53   1.61   
+ttHJet_ctWRe       -10.96  10.58  
+ttHJet_ctBRe       -12.86  13.11  
+ttHJet_cte1      -157.91 157.91 
+ttHJet_cte2      -157.91 157.91 
+ttHJet_cte3      -157.91 157.91 
+ttHJet_cleQt1Re1     -157.91 157.91 
+ttHJet_cleQt1Re2     -157.91 157.91 
+ttHJet_cleQt1Re3     -157.91 157.91 
+ttHJet_cleQt3Re1     -157.91 157.91 
+ttHJet_cleQt3Re2     -157.91 157.91 
+ttHJet_cleQt3Re3     -157.91 157.91 
+ttHJet_ctl      -157.91 157.91 
+ttHJet_ctHRe       -20.06  52.09  
+ttbar_cQe1    -157.91 157.91 
+ttbar_cQe2    -157.91 157.91 
+ttbar_cQe3    -157.91 157.91 
+ttbar_cQl31   -157.91 157.91 
+ttbar_cQl32   -157.91 157.91 
+ttbar_cQl33   -157.91 157.91 
+ttbar_cQl11   -157.91 157.91 
+ttbar_cQl12   -157.91 157.91 
+ttbar_cQl13   -157.91 157.91 
+ttbar_cbWRe     -157.91 157.91 
+ttbar_cHQ3    -157.91 157.91 
+ttbar_cHQ1    -157.91 157.91 
+ttbar_cHtb     -157.91 157.91 
+ttbar_cHtbRe    -157.91 157.91 
+ttbar_ctG     -13.58  6.92   
+ttbar_ctWRe     -56.29  55.34  
+ttbar_ctBRe     -66.05  66.81  
+ttbar_cte1    -157.91 157.91 
+ttbar_cte2    -157.91 157.91 
+ttbar_cte3    -157.91 157.91 
+ttbar_cleQt1Re1   -157.91 157.91 
+ttbar_cleQt1Re2   -157.91 157.91 
+ttbar_cleQt1Re3   -157.91 157.91 
+ttbar_cleQt3Re1   -157.91 157.91 
+ttbar_cleQt3Re2   -157.91 157.91 
+ttbar_cleQt3Re3   -157.91 157.91 
+ttbar_ctl    -157.91 157.91 
+ttbar_ctHRe     -157.91 157.91 
+ttllNuNuJetNoHiggs_cQe1     -43.53  43.91  
+ttllNuNuJetNoHiggs_cQe2     -43.53  43.91  
+ttllNuNuJetNoHiggs_cQe3     -43.53  43.91  
+ttllNuNuJetNoHiggs_cQl31    -157.91 157.91 
+ttllNuNuJetNoHiggs_cQl32    -157.91 157.91 
+ttllNuNuJetNoHiggs_cQl33    -157.91 157.91 
+ttllNuNuJetNoHiggs_cQl11    -42.99  44.86  
+ttllNuNuJetNoHiggs_cQl12    -42.99  44.86  
+ttllNuNuJetNoHiggs_cQl13    -42.99  44.86  
+ttllNuNuJetNoHiggs_cbWRe      -50.64  50.64  
+ttllNuNuJetNoHiggs_cHQ3     -54.68  52.25  
+ttllNuNuJetNoHiggs_cHQ1     -27.8   52.66  
+ttllNuNuJetNoHiggs_cHtb      -47.45  30.7   
+ttllNuNuJetNoHiggs_cHtbRe     -135.45 135.45 
+ttllNuNuJetNoHiggs_ctG      -3.76   2.62   
+ttllNuNuJetNoHiggs_ctWRe      -9.89   9.7    
+ttllNuNuJetNoHiggs_ctBRe      -6.95   7.01   
+ttllNuNuJetNoHiggs_cte1     -43.43  44.73  
+ttllNuNuJetNoHiggs_cte2     -43.43  44.73  
+ttllNuNuJetNoHiggs_cte3     -43.43  44.73  
+ttllNuNuJetNoHiggs_cleQt1Re1    -61.47  61.47  
+ttllNuNuJetNoHiggs_cleQt1Re2    -61.47  61.47  
+ttllNuNuJetNoHiggs_cleQt1Re3    -61.47  61.47  
+ttllNuNuJetNoHiggs_cleQt3Re1    -8.21   8.21   
+ttllNuNuJetNoHiggs_cleQt3Re2    -8.21   8.21   
+ttllNuNuJetNoHiggs_cleQt3Re2    -8.21   8.21   
+ttllNuNuJetNoHiggs_cleQt3Re3    -8.21   8.21   
+ttllNuNuJetNoHiggs_ctl     -43.53  44.26  
+ttllNuNuJetNoHiggs_ctHRe      -47.35  79.26  
+ttlnuJet_cQe1    -157.91 157.91 
+ttlnuJet_cQe2    -157.91 157.91 
+ttlnuJet_cQe3    -157.91 157.91 
+ttlnuJet_cQl31   -157.91 157.91 
+ttlnuJet_cQl32   -157.91 157.91 
+ttlnuJet_cQl33   -157.91 157.91 
+ttlnuJet_cQl11   -157.91 157.91 
+ttlnuJet_cQl12   -157.91 157.91 
+ttlnuJet_cQl13   -157.91 157.91 
+ttlnuJet_cbWRe     -157.91 157.91 
+ttlnuJet_cHQ3    -157.91 157.82 
+ttlnuJet_cHQ1    -157.91 157.91 
+ttlnuJet_cHtb     -157.91 157.91 
+ttlnuJet_cHtbRe    -157.91 157.91 
+ttlnuJet_ctG     -19.14  8.5    
+ttlnuJet_ctWRe     -12.73  12.16  
+ttlnuJet_ctBRe     -52.97  52.75  
+ttlnuJet_cte1    -157.91 157.91 
+ttlnuJet_cte2    -157.91 157.91 
+ttlnuJet_cte3    -157.91 157.91 
+ttlnuJet_cleQt1Re1   -157.91 157.91 
+ttlnuJet_cleQt1Re2   -157.91 157.91 
+ttlnuJet_cleQt1Re3   -157.91 157.91 
+ttlnuJet_cleQt3Re1   -157.91 157.91 
+ttlnuJet_cleQt3Re2   -157.91 157.91 
+ttlnuJet_cleQt3Re3   -157.91 157.91 
+ttlnuJet_ctl    -157.91 157.91 
+ttlnuJet_ctHRe     -157.91 157.91 

mcgeneration/configure_gridpack.py

@@ -4,7 +4,7 @@
 import random
 import time
 
-from helpers.helper_tools import linspace, parse_limit_file, find_process
+from helpers.helper_tools import linspace, parse_limit_file, find_process, make_restrict_card
 from helpers.ScanType import ScanType
 from helpers.BatchType import BatchType
 from helpers.DegreeOfFreedom import DegreeOfFreedom
@@ -258,7 +258,8 @@ def cmsconnect_chain_submit(gridpack,dofs,proc_list,tag_postfix,rwgt_pts,runs,st
             if tracker.getTarballTime(job) > 3*(tar_cut+delay):
                 # Skip checking jobs that finished sufficiently long ago
                 continue
-            if tracker.resubmitted.has_key(job) and tracker.resubmitted[job] >= resubmits:
+            # if tracker.resubmitted.has_key(job) and tracker.resubmitted[job] >= resubmits:
+            if job in tracker.resubmitted and tracker.resubmitted[job] >= resubmits:
                 # Stop trying to resubmit the job
                 continue
             p,c,r = job.split('_')
@@ -278,7 +279,8 @@ def cmsconnect_chain_submit(gridpack,dofs,proc_list,tag_postfix,rwgt_pts,runs,st
             gridpack.setProcess(p)
             #TODO: Might not want to split it up like this
             if stype == ScanType.SLINSPACE:
-                if proc_run_wl.has_key(p.getName()):
+                # if proc_run_wl.has_key(p.getName()):
+                if p.getName() in proc_run_wl:
                     submitted += submit_1dim_jobs(
                         gp=gridpack,
                         dofs=dofs,
@@ -355,7 +357,7 @@ def cmsconnect_chain_submit(gridpack,dofs,proc_list,tag_postfix,rwgt_pts,runs,st
 def submit_1dim_jobs(gp,dofs,npts,runs,tag_postfix='',max_submits=-1,run_wl={}):
     submitted = 0
     delay    =  10.0   # Time between successful submits (in seconds)
-    wc_limits = parse_limit_file(os.path.join("addons/limits","dim6top_LO_UFO_limits.txt"))
+    wc_limits = parse_limit_file(os.path.join("addons/limits","SMEFTsim_top_limits.txt"))
     for dof in dofs:
         dof_name = dof.getName()
         lim_key = "{process}_{wc}".format(process=gp.getOption('process'),wc=dof_name)
@@ -364,14 +366,50 @@ def submit_1dim_jobs(gp,dofs,npts,runs,tag_postfix='',max_submits=-1,run_wl={}):
             # The dof already has limits, re-use them
             low_lim = dof.getLow()
             high_lim = dof.getHigh()
-        elif wc_limits.has_key(lim_key):
+        # elif wc_limits.has_key(lim_key):
+        elif lim_key in wc_limits:
             # Use limits from the limits file for this process
             low_lim  = round(wc_limits[lim_key][0],6)
             high_lim = round(wc_limits[lim_key][1],6)
         else:
             low_lim,high_lim = gp.getOption('default_limits')
+
+        ###############################
+        restrict_ops = gp.getOption('restrict')
+        if restrict_ops:
+            # The getModel() method takes care of the case where we've specified a different model than
+            #   what the MGProcess process card uses
+            model = gp.getModel()
+            model_dir = "addons/models/{model}".format(model=model)
+            ref_restrict = restrict_ops['ref']
+            # We can't use str.removesuffix since it isn't available in python3.6
+            if ref_restrict.endswith('.dat'):
+                new_restrict = "{ref}_{name}".format(ref=ref_restrict.replace('.dat',''),name=dof.getName())
+            # The model import syntax in MG expects a specific naming convention -> MODEL-massless_NAME,
+            #   where _NAME corresponds to the restrict .dat card and would look like: restrict_massless_NAME.dat
+            new_model = "{model}-massless_{restrict}".format(model=model,restrict=dof.getName())
+            new_restrict = new_restrict + ".dat"
+
+            ref_restrict = os.path.join(model_dir,ref_restrict)
+            new_restrict = os.path.join(model_dir,new_restrict)
+            # Create a new restrict card in the directory of the model specified by the MGProcess
+            keep = restrict_ops['keep']
+            # Its ok to specify a parameter that doesn't show up in a particular lhablock. Also,
+            #   for operators that have been defined as linear combinations, we need to use the actual
+            #   parameter names, rather than what we get from dof.getName()
+            blocks = {k: list(dof.getCoefficients()) for k in restrict_ops['blocks']}
+            make_restrict_card(ref_restrict,new_restrict,keep=keep,**blocks)
+            # We can't use the 'replace_model' option to set the restrict card as this will overwrite
+            #   that option, which might have already been used to specify a different model. Modifying
+            #   the gridpack options like this feels rather dangerous, but not sure how to do this
+            #   any other way
+            restrict_ops['replace'] = [model,new_model]
+            gp.setOptions(restrict=restrict_ops)
+        ###############################
+
         for idx,start in enumerate(linspace(low_lim,high_lim,runs)):
-            if run_wl.has_key(dof_name) and idx not in run_wl[dof_name]:
+            # if run_wl.has_key(dof_name) and idx not in run_wl[dof_name]:
+            if dof_name in run_wl and idx not in run_wl[dof_name]:
                 continue
             pt = {}
             pt[dof.getName()] = start
@@ -457,7 +495,15 @@ def main():
     random.seed()
     stype = ScanType.FROMFILE
     btype = BatchType.CMSCONNECT
-    tag   = 'Run3_52WCs_SMEFTsim_top'
+    tag   = 'Example1'
+    # Set the restrict option to None or False to avoid using the restrict card machinary
+    restrict = {
+        "ref": "restrict_massless.dat", # Name of the restrict card to use as the reference
+        "blocks": ["SMEFT","SMEFTcpv"], # Name of the lhablock(s) that we want to modify
+        "keep": True,                   # Keep only the parameters we specify
+        "replace": None,                # This needs to be set prior to each Gridpack.setup() call
+    }
+    # restrict = None
     runs  = 1               # if set to 0, will only make a single gridpack
     npts  = 0
     #scan_files = [
@@ -569,6 +615,7 @@ def main():
     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"])
+    gridpack.setOptions(restrict=restrict)
     # For creating feynman diagrams
     #gridpack.setOptions(btype=BatchType.LOCAL,save_diagrams=True,replace_model="dim6top_LO_UFO_each_coupling_order_v2020-05-19")
     #gridpack.setOptions(coupling_string="FCNC=0 DIM6^2=1 DIM6_ctB^2=1 DIM6_ctW^2=1") # For example

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)
@@ -110,6 +111,35 @@ def setProcess(self,p):
             flavor_scheme=p.getFlavorScheme(self.CARD_DIR)
         )
 
+    # The (baseline) model is defined for the process card, which is specified by the MGProcess. If
+    #   the user has specified a new model via the 'replace_model' option, then the model we get from
+    #   the template process card likely won't be the same model that gets used in the actual process
+    #   card that MG reads. This is b/c we replace the 'import model' line in the >copied< version of
+    #   the process card.
+    def getModel(self):
+        if not self.getOption("process"):
+            raise RuntimeError("No process set")
+        if self.getOption("replace_model"):
+            # User has specified a different model to use, so we return that one
+            old, new = self.getOption("replace_model")
+            # If the new model actually corresponds to a restrict card, we need to only return the
+            #   piece that corresponds to the model as a whole. MG syntax for restrict cards is of
+            #   the form: MODEL-RESTRICT, this also means that we should never use a '-' anywhere in
+            #   the name of a model or a restrict card
+            new = new.split('-')[0]
+            return new
+        process_card = os.path.join(self.HOME_DIR,self.CARD_DIR,self.ops['template_dir'],self.ops['process_card'])
+        search_str = "import model"
+        with open(process_card,'r') as f:
+            for l in f.readlines():
+                # Ignore comments
+                l = l.split(sep="#",maxsplit=1)[0]
+                l = l.strip()
+                if l.startswith(search_str):
+                    model = l[len(search_str):]
+                    return model.strip()
+        raise RuntimeError("Unable to find 'import model' line in {pcard} for {pname}".format(pcard=self.ops['process_card'],pname=self.ops['process']))
+
     def loadRunCard(self):
         """
             Parses a MadGraph run card, which can then be modified independent
@@ -207,6 +237,20 @@ def saveProcessCard(self,indent=0):
             sed_str = "s|import model {old}|import model {new}|g".format(old=old,new=new)
             subprocess.Popen(['sed','-i','-e',sed_str,fpath]).communicate()
 
+        # Note: It is assumed that the user has made sure that the 'replace' info matches what the
+        #       current state of the copied process card is, e.g. if 'replace_model' has been used,
+        #       then the 'model' part of the 'replace' info should already correspond to w/e the
+        #       'replace_model' option changed things to.
+        # TODO: Since the getModel() method does take care to consider the case where the user has
+        #       specified the 'replace_model' option, we might be able to use that to help ensure
+        #       things are a little more consistent. The user would still be responsible for making
+        #       sure the restrict card exists and is placed in the correct location.
+        if self.ops['restrict']:
+            model, restrict = self.ops['restrict']['replace']
+            print("{ind}Using restrict card {restrict} for model {model}".format(ind=indent_str,model=model,restrict=restrict))
+            sed_str = "s|import model {model}|import model {restrict}|g".format(model=model,restrict=restrict)
+            subprocess.Popen(['sed','-i','-e',sed_str,fpath]).communicate()
+
         # Replace SUBSETUP in the process card with the correct name
         sed_str = "s|SUBSETUP|{setup}|g".format(setup=setup)
         subprocess.Popen(['sed','-i','-e',sed_str,fpath]).communicate()