generalising better

tracewin_errorstudy

+#!/usr/bin/env python
 
 
-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("-c", dest="calc_dir", default='temp', 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("-l", dest="lattice", default='lattice.dat', type=str, help="The lattice file for the project")
+    parser.add_argument("-t", dest="settings", default='settings.yml', type=str, help="The settings for the project template")
     parser.add_argument("-p", dest="priority", default=0, type=int, help="Job priority")
 
     return parser.parse_args()
@@ -32,34 +33,35 @@ def setup_exec_folder():
             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=''):
+def create_and_copy_lattice(args, 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
+    (note that a yaml template is required)
     '''
     import yaml,shutil,os
 
     # Create the lattice file from template:
-    lattice_settings=yaml.load(file(SETTINGS,'r').read())
+    lattice_settings=yaml.load(file(args.settings,'r').read())
     lattice_settings['extra_line']=extra_line
 
-    lattice_template=file(LATTICE+'.tmp','r').read()
+    lattice_template=file(args.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))
+    file(calc_dir+'/'+args.lattice,'w').write(lattice_template.format(**lattice_settings))
 
-    shutil.copy('settings.yml',calc_dir)
+    shutil.copy(args.settings, calc_dir)
 
     for folder in os.listdir(calc_dir):
         if folder[0:15]=="Local_TraceWin_":
-            path=os.path.join(calc_dir,folder,LATTICE)
+            path=os.path.join(calc_dir,folder,args.lattice)
             if os.path.exists(path):
                 os.remove(path)
-            os.link(calc_dir+'/'+LATTICE,path)
+            os.link(calc_dir+'/'+args.lattice,path)
 
 def create_extra_folder(calc_dir, f0, f1):
     import os
@@ -118,7 +120,7 @@ def setup_multi(args):
                 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')
+                create_and_copy_lattice(args, 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
 
@@ -179,7 +181,7 @@ def setup(args):
     # the statistical study in TW to do, normally 14
     stat_study=args.study+13
 
-    create_and_copy_lattice(args.calc_dir)
+    create_and_copy_lattice(args, args.calc_dir)
 
     if not os.path.isdir("TraceWin_exe"):
         setup_exec_folder()