deps/delphes_EIC/run_study.py

0001 #!/usr/bin/env python
0002 #SBATCH -J EICSTUDY    # job name
0003 #SBATCH -o logs/eicstudy-%A_%a.out
0004 #SBATCH -e logs/eicstudy-%A_%a.out
0005 #SBATCH -n 1
0006 #SBATCH -c 1
0007 #SBATCH --mem=6G
0008 #SBATCH -p htc             # queue (partition) -- batch, parallel, etc.  parallel-medium
0009 #SBATCH -t 12:00:00        # run time (hh:mm:ss)
0010 #SBATCH -D .               # Directory where executable will be run
0011 #SBATCH --mail-user=ssekula@smu.edu
0012 #SBATCH --mail-type=fail    # email me when the job finishes
0013
0014
0015 # This script was written originally for use on a SLURM-based batch system,
0016 # like the one available at SMU ("ManeFrame II"). It can be run on the command-line
0017 # alone; if the script doesn't detect the requested SLURM environment variables,
0018 # it will ask you to specify them. For instance, to run the first variation in a
0019 # study,
0020 #
0021 # SLURM_ARRAY_TASK_ID=0 ./run_study.py -t ./bfield_study_template.tcl -p "['PARAM_BFIELD']" -v "[[1.0,1.5,2.0,2.5,3.0,3.5]]" -n bfield
0022 #
0023 # The above will manually select the first variation in the array passed to "-v" and run it.
0024 #
0025 # General idea:
0026 #
0027 # A template script is executed by the Delphes simulation executable. Before execution,
0028 # placeholder parameters are filled in by the script and a copy of the template is
0029 # generated (with the parameter values set) for actual execution. For example:
0030 #
0031 # SLURM_ARRAY_TASK_ID=0 ./run_study.py -t ./bfield_study_template.tcl -p "['PARAM_BFIELD']" -v "[[2.0]]" -n bfield
0032 #
0033 # will open ./bfield_study_template.tcl, find all instanced of PARAM_BFIELD, and replace them with the number "2.0".
0034 # This allows command-line variation of one or more parameters to perform a study of changes to simulation.
0035 #
0036
0037 import subprocess
0038 import math
0039 import os
0040 import sys
0041 import shutil
0042 import glob
0043 import re
0044 import ast
0045
0046 import argparse
0047
0048 global tclparams, pythiaparams
0049 tclparams = {}
0050 pythiaparams = {}
0051
0052
0053 def LoadParamsTcl(tclfile):
0054     tcl = open(tclfile)
0055
0056     param_pattern = re.compile("^set\s+(.*?)\s+(.*)")
0057     for line in tcl:
0058         line = line.rstrip('\n')
0059         match = param_pattern.match(line)
0060         if match != None:
0061             tclparams[match.group(1)] = str(match.group(2))
0062
0063     print("Customization TCL Parameters are:")
0064     print(tclparams)
0065
0066 def LoadParamsPythia(cmndfile):
0067     cmnd = open(cmndfile)
0068
0069     param_pattern = re.compile("^\!\s*(.*?)\s*=\s*(.*)")
0070     for line in cmnd:
0071         line = line.rstrip('\n')
0072
0073         if line.find('END HEADER') != -1:
0074             break
0075
0076         match = param_pattern.match(line)
0077         if match != None:
0078             pythiaparams[match.group(1)] = str(match.group(2))
0079
0080     print("Customization Pythia8 Parameters are:")
0081     print(pythiaparams)
0082
0083
0084 def WriteTclFile(filename):
0085     global args
0086     tclfile = open(filename, 'w')
0087
0088     for param in args.params:
0089         value = args.params[param]
0090         if (param not in tclparams.keys()) and (param not in pythiaparams.keys()):
0091             print("WARNING: you tried to set %s, which is not an available parameter!" % (param))
0092             continue
0093         tclparams[param] = str(value)
0094
0095     for param in tclparams:
0096         line = "set %s %s\n" % (param, tclparams[param])
0097         tclfile.write(line)
0098     tclfile.close()
0099
0100 def WritePythiaFile(output_filename):
0101     global args
0102
0103     for param in args.params:
0104         value = args.params[param]
0105         if (param not in pythiaparams.keys()) and (param not in tclparams.keys()):
0106             print("WARNING: you tried to set %s, which is not an available parameter!" % (param))
0107             continue
0108         pythiaparams[param] = str(value)
0109
0110
0111     with open(args.commands, "rt") as input_template:
0112         with open(output_filename, "wt") as output_template:
0113             for line in input_template:
0114                 for param in pythiaparams:
0115                     value = str(pythiaparams[param])
0116                     line = line.replace(param, value)
0117
0118                 output_template.write(line)
0119
0120
0121 def TransferTclFile(tclfile, taskdir):
0122     with open(tclfile, "rt") as input_template:
0123         with open("%s/%s" % (taskdir, tclfile), "wt") as output_template:
0124             for line in input_template:
0125                 line = line.replace("customizations.tcl", "%s/customizations.tcl" % (taskdir))
0126
0127                 output_template.write(line)
0128
0129
0130
0131 parser = argparse.ArgumentParser()
0132
0133 parser.add_argument("-n", "--name", type=str,
0134                     help="name for this study set (e.g. bfield)")
0135 parser.add_argument("-t", "--template", type=str,
0136                     help="template TCL file for study")
0137 parser.add_argument("-c", "--commands", type=str,
0138                     help="template command file for study")
0139 parser.add_argument("-p", "--params", type=ast.literal_eval, default={},
0140                     help="environment variables to set for study")
0141 #parser.add_argument("-v", "--values", type=ast.literal_eval, default=[],
0142 #                    help="parameter values for study")
0143 parser.add_argument("-f", "--force", default=False, action='store_true',
0144                     help="force-overwrite existing output")
0145
0146 global args
0147 args = parser.parse_args()
0148
0149 # Create the task superdirectory
0150
0151 if not os.path.exists(args.name):
0152     try:
0153         os.makedirs(args.name)
0154     except OSError:
0155         print("%s already exists... continuing..." % (args.name))
0156
0157
0158 SLURM_ARRAY_TASK_ID="0"
0159
0160 try:
0161     SLURM_ARRAY_TASK_ID=os.environ["SLURM_ARRAY_TASK_ID"]
0162 except:
0163     print("Please set the SLURM_ARRAY_TASK_ID environment variable to a number (e.g. 0) before running this script.")
0164     sys.exit()
0165     pass
0166
0167 print("Task ID requested: %d" % (int(SLURM_ARRAY_TASK_ID)))
0168
0169
0170
0171 value_index = int(SLURM_ARRAY_TASK_ID)
0172
0173
0174 # Load the available parameters and their default values from
0175 # customizations.tcl and from the header of the PYTHIA8 command
0176 # file
0177
0178 LoadParamsTcl("customizations.tcl")
0179 LoadParamsPythia(args.commands)
0180
0181
0182 # Handle random number seed issues and other defaults
0183
0184 random_seed = 0
0185 if "PARAM_RANDOM_SEED" not in args.params:
0186     random_seed = abs(hash(args.name)) % (10 ** 8) + value_index
0187     pythiaparams["PARAM_RANDOM_SEED"] = random_seed
0188
0189 # Execute the study
0190
0191 taskdir="%s/%d" % (args.name, value_index)
0192 tclfile = "%s/customizations.tcl" % (taskdir)
0193 cmndfile = "%s/%s" % (taskdir, args.commands)
0194
0195 if os.path.exists(taskdir) and not args.force:
0196     print("Skipping this task directory --- it already exists. Cleanup before overwriting!")
0197     print(taskdir)
0198 else:
0199     if not os.path.exists(taskdir):
0200         os.makedirs(taskdir)
0201
0202     # Replace parameter placeholders with values for this study
0203     WriteTclFile(tclfile)
0204     WritePythiaFile(cmndfile)
0205
0206     # Copy files to the task directory before execution
0207     copy_files = [args.template]
0208     for a_file in copy_files:
0209         subprocess.call("cp %s %s" % (a_file, taskdir), shell=True)
0210         #TransferTclFile(a_file, taskdir)
0211     # Write the random number seed to disk
0212     rndm_file = open(taskdir+"/random_seed.dat", "w")
0213     rndm_file.write(str(random_seed))
0214     rndm_file.close()
0215
0216     # Execute the study
0217     subprocess.call("DelphesPythia8 {0[taskdir]}/{0[template]} {0[taskdir]}/{0[commands]} {0[taskdir]}/out.root".format({'taskdir': taskdir, 'template': args.template, 'commands': args.commands}), shell=True)