deps/delphes_EIC/run_skim.py

0001 #!/usr/bin/env python
0002 #SBATCH -J EICSKIM    # job name
0003 #SBATCH -o logs/eicskim-%A_%a.out
0004 #SBATCH -e logs/eicskim-%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 00:05: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_skim.py -i <INPUT DIRECTORY> -o <OUTPUT DIRECTORY> -f out.root -c "Jet.Flavor==4"
0022 #
0023 # will skim the <INPUT DIRECTORY>/0/out.root file into <OUTPUT DIRECTORY>/0/out.root file, using the cut
0024 # Jet.Flavor==4. Any event containing such a jet will be retained in the skim.
0025
0026 import subprocess
0027 import math
0028 import os
0029 import sys
0030 import shutil
0031 import glob
0032 import re
0033 import ast
0034
0035 import argparse
0036 parser = argparse.ArgumentParser()
0037
0038 parser.add_argument("-i", "--input", type=str,
0039                     help="directory holding all input ROOT files for skimming")
0040 parser.add_argument("-o", "--output", type=str,
0041                     help="directory holding all input ROOT files for skimming")
0042 parser.add_argument("-r", "--rootfile", type=str,
0043                     help="Name of the ROOT file in each subdirectory of the input directory")
0044 parser.add_argument("-c", "--cuts", type=str,
0045                     help="ROOT selection string-style cuts")
0046 parser.add_argument("-f", "--force", default=False, action='store_true',
0047                     help="force-overwrite existing output")
0048
0049 args = parser.parse_args()
0050
0051
0052
0053 # Create the task superdirectory
0054
0055 if not os.path.exists(args.output):
0056     try:
0057         os.makedirs(args.output)
0058     except OSError:
0059         print("%s already exists... continuing..." % (args.output))
0060
0061
0062 SLURM_ARRAY_TASK_ID="0"
0063
0064 try:
0065     SLURM_ARRAY_TASK_ID=os.environ["SLURM_ARRAY_TASK_ID"]
0066 except:
0067     print("Please set the SLURM_ARRAY_TASK_ID environment variable to a number (e.g. 0) before running this script.")
0068     sys.exit()
0069     pass
0070
0071 print("Task ID requested: %d" % (int(SLURM_ARRAY_TASK_ID)))
0072
0073
0074
0075 value_index = int(SLURM_ARRAY_TASK_ID)
0076
0077
0078 # Execute the skim
0079
0080 taskdir="%s/%d" % (args.output, value_index)
0081 inputfile="%s/%d/%s" % (args.input, value_index, args.rootfile)
0082 outputfile="%s/%d/%s" % (args.output, value_index, args.rootfile)
0083
0084 if os.path.exists(taskdir) and not args.force:
0085     print("Skipping this task directory --- it already exists. Cleanup before overwriting!")
0086     print(taskdir)
0087 else:
0088     if not os.path.exists(taskdir):
0089         os.makedirs(taskdir)
0090
0091     # Execute the study
0092     subprocess.call("root -q -l -b ./DelphesSkim.C'+(\"{0[input]}\",\"{0[output]}\",\"{0[cuts]}\")'".format({'input': inputfile, 'output': outputfile, 'cuts': args.cuts}), shell=True)