added tracewin error study script

starts ht condor run

tracewin_errorstudy

+
+
+SETTINGS='settings.yml'
+LATTICE='lattice.dat'
+
+def define_and_handle_args():
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Simple setup script to run TraceWin on HT Condor..")
+    parser.add_argument("-s", dest="seed", default=None, type=int, help="Define seed")
+    parser.add_argument("-c", dest="calc_dir", default='backup_nosync', type=str, help="Path to calculation folder")
+    parser.add_argument("-n", dest="num_jobs", default=0, type=int, help="Number of jobs")
+    parser.add_argument("-t", dest="study", default=1, type=int, help="Statistical study number")
+    parser.add_argument("-m", dest="multi", default=1, type=int, help="Multi-dynamic study")
+    parser.add_argument("-p", dest="priority", default=0, type=int, help="Job priority")
+
+    return parser.parse_args()
+
+def setup_exec_folder():
+    import os,shutil
+
+    orig_path=os.path.join(os.environ.get("HOME"),"TraceWin_exe")
+    if not os.path.isdir(orig_path):
+        raise ValueError("Cannot find tracewin executables in "+orig_path+" yet")
+    os.makedirs("TraceWin_exe")
+
+    o=dict(LINUX="lx",OSX="mac")
+    a=dict(INTEL="",X86_64="64")
+
+    for OpSys in ["LINUX", "OSX"]:
+        for Arch in ["INTEL", "X86_64"]:
+            shutil.copy(os.path.join(orig_path,"trace"+o[OpSys]+a[Arch]),
+                        os.path.join("TraceWin_exe","tracewin.%s.%s" % (OpSys,Arch)))
+
+def create_and_copy_lattice(calc_dir, commented=[], extra_line=''):
+    '''
+    Creates the lattice file from tmp
+    comments out any commands listed in commented
+    
+    For multi-study set the extra_line
+    '''
+    import yaml,shutil,os
+
+    # Create the lattice file from template:
+    lattice_settings=yaml.load(file(SETTINGS,'r').read())
+    lattice_settings['extra_line']=extra_line
+
+    lattice_template=file(LATTICE+'.tmp','r').read()
+    for cmd in commented:
+        lattice_template=lattice_template.replace(cmd.lower(),';'+cmd)
+        lattice_template=lattice_template.replace(cmd.upper(),';'+cmd)
+
+    file(calc_dir+'/'+LATTICE,'w').write(lattice_template.format(**lattice_settings))
+
+    shutil.copy('settings.yml',calc_dir)
+
+    for folder in os.listdir(calc_dir):
+        if folder[0:15]=="Local_TraceWin_":
+            path=os.path.join(calc_dir,folder,LATTICE)
+            if os.path.exists(path):
+                os.remove(path)
+            os.link(calc_dir+'/'+LATTICE,path)
+
+def create_extra_folder(calc_dir, f0, f1):
+    import os
+
+    path0=os.path.join(calc_dir,f0)
+    path1=os.path.join(calc_dir,f1)
+
+    if not os.path.isdir(path1):
+        os.makedirs(path1)
+
+    for f in os.listdir(path0):
+        p0=os.path.join(path0,f)
+        if f[-2:]=='_0':
+            p1=os.path.join(path1,f[:-1])+path1.split('_')[-1]
+        else:
+            p1=os.path.join(path1,f)
+        if not os.path.exists(p1):
+            # hard link
+            os.link(p0,p1)
+
+def delete_output_files(args,folder):
+    '''
+    Remove all output files in a folder 
+    (clean it for a re-run)
+
+    '''
+    import os
+
+    output_files=['ENV_diag1.dat', 'part_dtl1.dst', 'part_rfq.dst', 'Remote_DATA.dat', 'condor.log', 'dtl1.plt', 'PAR_diag1.dat', 'Density_PAR.dat_9', 'lattice.par']
+
+    full_path=os.path.join(args.calc_dir,folder)
+    for f in os.listdir(full_path):
+        for of in output_files:
+            if of in f:
+                os.remove(os.path.join(full_path,f))
+
+def setup_multi(args):
+    import os
+
+    print "multi study.."
+
+    # the statistical study in TW to do, normally 14
+    stat_study=args.study+13
+    # for multi study we need an extra line to set dynamic errors:
+
+    num_jobs=0
+    new_folders=[]
+    for njob in xrange(args.num_jobs):
+        folder='Local_TraceWin_{}'.format(njob)
+        if os.path.isfile(os.path.join(args.calc_dir,folder,'Error_Datas.txt')):
+            for mjob in xrange(args.multi):
+                num_jobs+=1
+                if mjob==int(folder.split('_')[-1]):
+                    continue
+                newfolder='{}_{}'.format(folder,mjob)
+                new_folders.append(newfolder)
+                create_extra_folder(args.calc_dir,folder,newfolder)
+                delete_output_files(args,newfolder)
+                create_and_copy_lattice(os.path.join(args.calc_dir,newfolder),commented=['ERROR_CAV_CPL_DYN', 'ERROR_CAV_NCPL_DYN'],extra_line='ERROR_DYN_FILE Error_Datas.txt')
+
+    print new_folders
+
+    if not os.path.isdir("TraceWin_exe"):
+        setup_exec_folder()
+
+    # Get random seeds from previously run jobs:
+    seeds=[]
+    for l in file(os.path.join(args.calc_dir,'submit.job')):
+        if len(l.strip()) and l.split()[0]=='arguments':
+            seeds.append(l.split()[-1])
+
+
+    head_template=file("head.job.multi.tmp","r").read()
+    queue_template=file("queue.job.multi.tmp","r").read()
+    with file(args.calc_dir+"/multi.job","w") as out:
+        out.write(head_template.format(**locals()))
+
+        for job_folder in new_folders:
+            i=int(job_folder.split('_')[-2])
+            j=int(job_folder.split('_')[-1])
+
+            # Get file list:
+            input_files=os.listdir(os.path.join(args.calc_dir,job_folder))
+            input_files=', '.join(input_files)
+
+            q_priority=args.priority
+            rand=seeds[j]
+            out.write(queue_template.format(**locals()))
+
+    if num_jobs==0:
+        raise ValueError("Could not find any simulation folders in "+args.calc_dir)
+
+    print "Total number of new jobs", num_jobs
+
+def setup(args):
+    import os,random
+
+    if args.multi>1:
+        setup_multi(args)
+        return 0
+
+    num_jobs=0
+    for folder in os.listdir(args.calc_dir):
+        if folder[0:15]=="Local_TraceWin_":
+            num_jobs+=1
+
+    if num_jobs==0:
+        raise ValueError("Could not find any simulation folders in "+args.calc_dir)
+
+    # Create extra folders if needed..
+    if args.num_jobs>num_jobs:
+        for n in xrange(num_jobs,args.num_jobs):
+            folder_name='Local_TraceWin_{n}'.format(**locals())
+            create_extra_folder(args.calc_dir,'Local_TraceWin_0',folder_name)
+        num_jobs=args.num_jobs
+
+    # the statistical study in TW to do, normally 14
+    stat_study=args.study+13
+
+    create_and_copy_lattice(args.calc_dir)
+
+    if not os.path.isdir("TraceWin_exe"):
+        setup_exec_folder()
+
+    if args.seed!=None:
+        random.seed(args.seed)
+    rand_lst=[random.randint(1e5,1e9) for i in xrange(num_jobs)]
+    #rand_lst=file('rand.dat','r').read().split()
+
+    # Get file list:
+    input_files=os.listdir(os.path.join(args.calc_dir,"Local_TraceWin_0"))
+    for i in range(len(input_files)):
+        if input_files[i][-2:]=='_0':
+            input_files[i]=input_files[i][:-2]+'_$(Process)'
+    input_files=', '.join(input_files)
+
+    head_template=file("head.job.tmp","r").read()
+    queue_template=file("queue.job.tmp","r").read()
+    with file(args.calc_dir+"/submit.job","w") as out:
+        out.write(head_template.format(**locals()))
+
+        for i in xrange(num_jobs):
+            q_priority=max(-5,args.priority-i/100)
+            rand=rand_lst[i]
+            out.write(queue_template.format(**locals()))
+
+
+def execute(calc_dir,multi):
+    import subprocess
+    if multi:
+        cmd=["condor_submit",calc_dir+"/multi.job"]
+    else:
+        cmd=["condor_submit",calc_dir+"/submit.job"]
+    subprocess.call(cmd)
+
+if __name__=="__main__":
+    args=define_and_handle_args()
+    setup(args)
+    execute(args.calc_dir,args.multi>1)