Something went wrong on our end
Code owners
Assign users and groups as approvers for specific file changes. Learn more.
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)