import random
from datetime import datetime
import os
import subprocess
import inspect
import argparse
import sys
import re
def parsestat(s):
return int(s.
replace("k","000").
replace("M","000000").
replace("G","000000000")
)
def getref():
try:
rev = subprocess.check_output([
'git', 'rev-parse', '--short', 'HEAD'
], stderr=subprocess.PIPE).splitlines()[0]
return rev
except subprocess.CalledProcessError:
return "not on git"
def getbranch():
try:
rev = subprocess.check_output([
'git', 'rev-parse', '--abbrev-ref', 'HEAD'
], stderr=subprocess.PIPE).splitlines()[0]
return rev
except subprocess.CalledProcessError:
return "not on git"
def rerunconfig(
flavour="mu-e", whichbase="m2enn",
seeds=5, xicuts=[.5, .25, .125],
stats={
"0": (10000, 20, 100000, 100),
"V": (10000, 20, 100000, 100),
"R": (20000, 20, 500000, 100)
}, folder="", binary='xs_main'
):
if len(folder) == 0:
folder = whichbase + flavour
cargs = [
"--seeds %s" % ' '.join(str(i) for i in seeds),
"-xi %s" % ' '.join(str(i) for i in xicuts),
"--flavour %s" % flavour,
"--piece %s" % whichbase,
"--output-dir %s" % folder,
"--prog %s" % binary
]
for part, stat in stats.iteritems():
if len(stat) == 5:
cargs.append("--stat %s,%d,%d,%d,%d,%d" % (
part, stat[0], stat[1], stat[2], stat[3], stat[4])
)
else:
cargs.append("--stat %s,%d,%d,%d,%d" % (
part, stat[0], stat[1], stat[2], stat[3])
)
return cargs
def create_menu(
flavour="mu-e", whichbase="m2enn",
seeds=5, xicuts=[.5, .25, .125],
stats={
"0": (10000, 20, 100000, 100),
"V": (10000, 20, 100000, 100),
"R": (20000, 20, 500000, 100)
}, folder="", binary='xs_main'
):
totalstat = [0,0]
if type(seeds) == int:
n = seeds # Get n random numbers
seeds = []
while len(seeds) < n:
r = random.randint(10000, 100000)
if r in seeds:
continue
seeds.append(r)
elif type(seeds) == list:
pass
else:
raise TypeError(
"Seeds needs to be either list or int, is type "
+ str(type(seeds))
)
if len(folder) == 0:
folder = whichbase + flavour
rrc = rerunconfig(
flavour, whichbase, seeds, xicuts,
stats, folder, binary
)
print "Building files. To rerun this, execute"
print "python %s \\" % sys.argv[0]
print " \\\n".join(
" " + i for i in rrc
)
header = (
"## Generated at %s by %s\n"
+ "# git version: %s (%s)\n"
+ "# To re-generate, run python %s \\\n"
) % (
datetime.now().strftime("%H:%M on %B %d %Y"), os.environ['USER'],
getbranch(), getref(), sys.argv[0]
) + " \\\n".join(
"# " + i for i in rrc
)
menu = "\n\nconf %s/%s-%s.conf" % (folder, whichbase, flavour)
config = ("\n\n# specify the program to run relative to `pwd`\n"
+ "binary=%s\n\n"
+ "# specify the output folder\n"
+ "folder=%s/\n\n"
+ "# Specify the variables nenter_ad, itmx_ad, nenter and itmx\n"
+ "# for each piece you want to run.\n"
+ "declare -A STAT=(\n") % (binary, folder)
for part,stat in stats.iteritems():
config += ' ["%s%s"]="%d\\n%d\\n%d\\n%d"\n' % (
whichbase, part, stat[0], stat[1], stat[2], stat[3]
)
if stat[0] < stat[1] or stat[2] < stat[3]:
print "Warning! In part %s, calls and itmx might be swapped!" % part
n = stat[4] if len(stat) == 5 else len(seeds)
xiloop = (
"R" in part or # real corrections
"F" in part or # real-virtual corrections
"CT" in part or # explicit counter term
("V" in part and # virtual corrections and ...
"CT" not in "".join(stats.keys())) # ... never and explicit CT
)
for xi in (xicuts if xiloop else [1.]):
totalstat[0] += n*(stat[1] + stat[3]) # iterations
totalstat[1] += 1000*n*(stat[0]*stat[1] + stat[2]*stat[3]) # iterations
menu += "\n\n"
menu += '\n'.join(
"run %d %f %s%s %s 0" % (seed, xi, whichbase, part, flavour)
for seed in seeds[:n]
)
menu += "\n\n"
config += ")"
if totalstat[1] > 999999999999:
c = 'T'
p = float(totalstat[1])/1.e12
elif totalstat[1] > 999999999:
c = 'G'
p = float(totalstat[1])/1.e9
elif totalstat[1] > 999999:
c = 'M'
p = float(totalstat[1])/1.e6
elif totalstat[1] > 999:
c = 'k'
p = float(totalstat[1])/1.e3
print "Expect %d iterations, %f%s calls" % (
totalstat[0], p,c
)
return (
["%s/menu-%s-%s.txt" % (folder, whichbase, flavour), header + menu],
["%s/%s-%s.conf" % (folder, whichbase, flavour), header + config],
folder
)
def interogate(args={}):
specs = inspect.getargspec(create_menu)
defaults = dict(zip(specs.args, specs.defaults))
def basic_ask(msg, default, parser=(lambda x: x)):
ans = raw_input("%s? [%s] " % (msg, default))
if ans == '':
return default
else:
return parser(ans)
def ask(key, msg, parser=(lambda x: x)):
if key not in args:
args[key] = basic_ask(msg, defaults[key], parser)
ask('whichbase', "What type of process")
ask('flavour', "Which flavour combination")
ask('seeds', "How many / which seeds",
parser=lambda x: int(x) if ',' not in x else [int(i) for i in x.split(',')]
)
ask('xicuts', "Which xi cuts",
parser=lambda x: [float(i) for i in x.split(',')]
)
defaults['folder'] = args['whichbase'] + args['flavour']
ask('folder', "Where to store data")
ask('binary', "Which binary to run")
if 'stats' not in args:
pieces = basic_ask(
"Which pieces", ['0', 'V', 'R'],
parser=lambda x: x.split(',')
)
args['stats'] = {}
for piece in pieces:
args['stats'][piece] = basic_ask(
"How much statistics for " + piece + " (pc, pi, c, i)",
(10000, 20, 100000, 100),
parser=lambda x: tuple(parsestat(i) for i in x.split(','))
)
return args
def parseargs():
specs = inspect.getargspec(create_menu)
defaults = dict(zip(specs.args, specs.defaults))
parser = argparse.ArgumentParser(
description='Generate menu and config files'
)
parser.add_argument(
'-i', action='store_true',
help='reads all arguments and questions the user about the rest'
)
group = parser.add_mutually_exclusive_group()
group.add_argument(
'-ns', '--nseeds', type=int,
help="number of random seeds to use [%d]" % defaults['seeds']
)
group.add_argument('--seeds', nargs='+', type=int, help="list of seeds")
parser.add_argument(
'-xi', nargs='+', type=float,
help="list of xicut [%s]" % (
','.join(str(i) for i in defaults['xicuts'])
)
)
parser.add_argument(
'-f', '--flavour', type=str,
help="flavour variable [%s]" % defaults['flavour']
)
parser.add_argument(
'-p', '--piece', type=str,
help="the which_piece prefix [%s]" % defaults['whichbase']
)
parser.add_argument(
'-d', '--output-dir', help="output folder, aborts if folder exists"
)
parser.add_argument(
'--force', action='store_true', help="overwrite existing files"
)
parser.add_argument(
'--prog', default='xs_main', help="executable of the MC [xs_main]"
)
parser.add_argument(
'-s', '--stat', action="append",
help="comma seperated list of piece,pc,pi,c,i[,seeds to use]"
)
parsed = parser.parse_args()
args = {}
if parsed.nseeds:
args['seeds'] = parsed.nseeds
if parsed.seeds:
args['seeds'] = parsed.seeds
if parsed.xi:
args['xicuts'] = parsed.xi
if parsed.stat:
args['stats'] = {}
for stat in parsed.stat:
s = stat.split(',')
args['stats'][s[0]] = tuple(parsestat(i) for i in s[1:])
elif not parsed.i:
args['stats'] = defaults['stats']
if parsed.flavour:
args['flavour'] = parsed.flavour
if parsed.piece:
args['whichbase'] = parsed.piece
if parsed.output_dir:
args['folder'] = parsed.output_dir
if parsed.prog:
args['binary'] = parsed.prog
if parsed.i:
args = interogate(args)
elif not parsed.output_dir:
parser.error("Check your flags")
return args, parsed.force
def save(menu, conf, folder, force=False):
toolspath = os.path.dirname(os.path.realpath(__file__))
if os.path.isdir(folder):
print "The folder ", folder, " already exists."
if not force:
raise SystemExit("File exists")
else:
print "Overwritting."
else:
os.mkdir(folder)
os.mkdir(folder+"/out/")
if 'afs' in toolspath:
print "You are running on AFS, won't fix permissions"
else:
os.chmod(folder, 0o2775)
os.chmod(folder+"/out/", 0o2775)
with open(toolspath + "/submit.sh") as fp:
submit = fp.read()
submit = re.sub(
"^#SBATCH --output=.*/slurm-%j.out$",
"#SBATCH --output=%s/slurm-%%j.out" % folder,
submit, flags=re.M
)
submit = re.sub(
"^#SBATCH --input=.*$",
"#SBATCH --input=%s" % menu[0],
submit, flags=re.M
)
submit = re.sub(
"^this=\./submit\.sh$",
"this=%s/submit.sh" % folder,
submit, flags=re.M
)
with open(menu[0], 'w') as fp:
fp.write(menu[1])
with open(conf[0], 'w') as fp:
fp.write(conf[1])
with open(folder + "/submit.sh", 'w') as fp:
fp.write(submit)
os.chmod(menu[0], 0o664)
os.chmod(folder + "/submit.sh", 0o775)
print "Created menu, config and submit script in ",folder
print "Please change the ntasks and time options accordingly"
if __name__ == '__main__':
import sys
args, force = parseargs()
menu, config, folder = create_menu(**args)
save(menu, config, folder, force=force)