MillePedeAlignmentAlgorithm/scripts/mps_splice.py

0001 #!/usr/bin/env python3
0002 from builtins import range
0003 import re
0004 import argparse
0005 import math
0006 import fileinput
0007
0008 # Set up argrument parser
0009 helpEpilog = ''''''
0010
0011 parser = argparse.ArgumentParser(description='Take card file, blank all INFI directives and insert the INFI directives from the modifier file instead.',
0012                                  epilog=helpEpilog,
0013                                  formatter_class=argparse.RawDescriptionHelpFormatter)
0014
0015 parser.add_argument('inCfg', action='store',
0016                     help='name of the config-template')
0017 parser.add_argument('modCfg', action='store',
0018                     help='name of the modifier file from mps_split')
0019 parser.add_argument('outCfg', action='store',
0020                     help='name of modified output file')
0021 parser.add_argument('isn', action='store',
0022                     help='number of the job (three digit number with leading zeros)')
0023 parser.add_argument("--max-events", dest = "max_events", type = int,
0024                     help = "maximum number of events to process")
0025 parser.add_argument("--skip-events", dest = "skip_events", type = int,
0026                     help = "number of events to skip before processing")
0027
0028 # Parse arguments
0029 args = parser.parse_args()
0030 inCfg  = args.inCfg
0031 modCfg = args.modCfg
0032 outCfg = args.outCfg
0033 isn    = args.isn
0034
0035
0036 # open input file
0037 with open(inCfg, 'r') as INFILE:
0038     body = INFILE.read()
0039
0040 # read modifier file
0041 with open(modCfg, 'r') as MODFILE:
0042     mods = MODFILE.read()
0043 mods = mods.strip()
0044
0045 # prepare the new fileNames directive. Delete first line if necessary.
0046 fileNames = mods.split('\n')
0047 if 'CastorPool=' in fileNames[0]:
0048     del fileNames[0]
0049
0050 # replace ISN number (input is a string of three digit number with leading zeros though)
0051 body = re.sub(re.compile('ISN',re.M), isn, body)
0052
0053 # print to outCfg
0054 with open(outCfg, 'w') as OUTFILE:
0055     OUTFILE.write(body)
0056
0057 # Number of total files and number of extends for cms.vstring are needed
0058 numberOfFiles   = len(fileNames)
0059 numberOfExtends = int(math.ceil(numberOfFiles/255.))
0060
0061 # Create and insert the readFile.extend lines
0062 for j in range(numberOfExtends):
0063     insertBlock = "readFiles.extend([\n    "
0064     i=0
0065     currentStart = j*255
0066     while (i<255) and ((currentStart+i)<numberOfFiles):
0067         entry = fileNames[currentStart+i].strip()
0068         if (i==254) or ((currentStart+i+1)==numberOfFiles):
0069             insertBlock += "\'"+entry+"\'])\n"
0070         else:
0071             insertBlock += "\'"+entry+"\',\n    "
0072         i+=1
0073
0074     for line in fileinput.input(outCfg, inplace=1):
0075         print(line,end='')
0076         if re.match('readFiles\s*=\s*cms.untracked.vstring()',line):
0077             print(insertBlock,end='')
0078
0079 if args.skip_events is not None:
0080     with open(outCfg, "a") as f:
0081         f.write("process.source.skipEvents = cms.untracked.uint32({0:d})\n"
0082                 .format(args.skip_events))
0083
0084 if args.max_events is not None:
0085     with open(outCfg, "a") as f:
0086         f.write("process.maxEvents = cms.untracked.PSet(input = "
0087                 "cms.untracked.int32({0:d}))\n".format(args.max_events))