made scripts executable to allow pre-commit to detect

updated tracewin_errorstudy to python3
corrected errors from black/flake8

scripts/tracewin_errorstudy

@@ -2,106 +2,128 @@
 from __future__ import print_function
 from ess import __queue_templates__
 
+
 def define_and_handle_args():
     import argparse
 
     parser = argparse.ArgumentParser(description="Simple setup script to run TraceWin on HT Condor.. Default values in square brackets")
     parser.add_argument("-s", dest="seed", default=None, type=int, help="Define seed [None]")
-    parser.add_argument("-c", dest="calc_dir", default='temp', type=str, help="Path to calculation folder [temp/]")
+    parser.add_argument("-c", dest="calc_dir", default="temp", type=str, help="Path to calculation folder [temp/]")
     parser.add_argument("-n", dest="num_jobs", default=0, type=int, help="Number of jobs [Count number of copied runs in TW]")
     parser.add_argument("-t", dest="study", default=1, type=int, help="Statistical study number [1]")
     parser.add_argument("-m", dest="multi", default=1, type=int, help="Multi-dynamic study [1]")
-    parser.add_argument("-l", dest="lattice", default='lattice.dat', type=str, help="The lattice file for the project [lattice.dat]")
-    parser.add_argument("-q", dest="settings", default='settings.yml', type=str, help="The settings for the project template (yaml) [settings.yml]")
+    parser.add_argument("-l", dest="lattice", default="lattice.dat", type=str, help="The lattice file for the project [lattice.dat]")
+    parser.add_argument(
+        "-q", dest="settings", default="settings.yml", type=str, help="The settings for the project template (yaml) [settings.yml]"
+    )
     parser.add_argument("-p", dest="priority", default=0, type=int, help="Job priority")
 
     return parser.parse_args()
 
+
 def setup_exec_folder():
-    import os,shutil
+    import os
+    import shutil
 
-    orig_path=os.path.join(os.environ.get("HOME"),"TraceWin/exe")
+    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")
+        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")
+    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)))
+            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(args, 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
+    """
+    import yaml
+    import shutil
+    import os
 
     if os.path.exists(args.settings):
         # Create the lattice file from template:
-        lattice_settings=yaml.load(open(args.settings,'r').read())
-        lattice_settings['extra_line']=extra_line
+        lattice_settings = yaml.load(open(args.settings, "r").read())
+        lattice_settings["extra_line"] = extra_line
 
-        lattice_template=open(args.lattice+'.tmp','r').read()
+        lattice_template = open(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)
+            lattice_template = lattice_template.replace(cmd.lower(), ";" + cmd)
+            lattice_template = lattice_template.replace(cmd.upper(), ";" + cmd)
 
-        open(os.path.join(calc_dir,args.lattice),'w').write(lattice_template.format(**lattice_settings))
+        open(os.path.join(calc_dir, args.lattice), "w").write(lattice_template.format(**lattice_settings))
 
         shutil.copy(args.settings, calc_dir)
 
     else:
-        if calc_dir != '.':
+        if calc_dir != ".":
             shutil.copy(args.lattice, os.path.join(calc_dir, args.lattice))
 
     for folder in os.listdir(calc_dir):
-        if folder[0:15]=="Local_TraceWin_":
-            path=os.path.join(calc_dir,folder,args.lattice)
+        if folder[0:15] == "Local_TraceWin_":
+            path = os.path.join(calc_dir, folder, args.lattice)
             if os.path.exists(path):
                 os.remove(path)
-            os.link(calc_dir+'/'+args.lattice,path)
+            os.link(calc_dir + "/" + args.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)
+    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]
+        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)
+            p1 = os.path.join(path1, f)
         if not os.path.exists(p1):
             # hard link
-            os.link(p0,p1)
+            os.link(p0, p1)
+
 
-def delete_output_files(args,folder):
-    '''
-    Remove all output files in a folder 
+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)
+    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))
+                os.remove(os.path.join(full_path, f))
+
 
 def setup_multi(args):
     import os
@@ -109,131 +131,141 @@ def setup_multi(args):
     print("multi study..")
 
     # the statistical study in TW to do, normally 14
-    stat_study=args.study+13
+    stat_study = args.study + 13
     # for multi study we need an extra line to set dynamic errors:
 
-    num_jobs=0
-    new_folders=[]
+    num_jobs = 0
+    new_folders = []
     for njob in range(args.num_jobs):
-        folder='Local_TraceWin_{}'.format(njob)
-        if os.path.isfile( os.path.join(args.calc_dir, folder, 'Error_Datas.txt')):
+        folder = "Local_TraceWin_{}".format(njob)
+        if os.path.isfile(os.path.join(args.calc_dir, folder, "Error_Datas.txt")):
             for mjob in range(args.multi):
-                num_jobs+=1
-                if mjob==int(folder.split('_')[-1]):
+                num_jobs += 1
+                if mjob == int(folder.split("_")[-1]):
                     continue
-                newfolder='{}_{}'.format(folder,mjob)
+                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(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')
+                create_extra_folder(args.calc_dir, folder, newfolder)
+                delete_output_files(args, newfolder)
+                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",
+                )
 
     if not os.path.isdir("TraceWin/exe"):
         setup_exec_folder()
 
     # Get random seeds from previously run jobs:
-    seeds=[]
-    for l in open(os.path.join(args.calc_dir,'submit.job'), 'r'):
-        if len(l.strip()) and l.split()[0]=='arguments':
+    seeds = []
+    for l in open(os.path.join(args.calc_dir, "submit.job"), "r"):
+        if len(l.strip()) and l.split()[0] == "arguments":
             seeds.append(l.split()[-1])
 
-
-    head_template=__queue_templates__.get_head_multi_template()
-    queue_template=__queue_templates__.get_queue_multi_template()
-    with open(args.calc_dir+"/multi.job","w") as out:
+    head_template = __queue_templates__.get_head_multi_template()
+    queue_template = __queue_templates__.get_queue_multi_template()
+    with open(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])
+            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)
+            input_files = os.listdir(os.path.join(args.calc_dir, job_folder))
+            input_files = ", ".join(input_files)
 
-            q_priority=args.priority
-            rand=seeds[j]
+            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)
+    if num_jobs == 0:
+        raise ValueError("Could not find any simulation folders in " + args.calc_dir)
 
     print("Total number of new jobs", num_jobs)
-    args.num_jobs=num_jobs
+    args.num_jobs = num_jobs
+
 
 def setup(args):
-    import os,random
+    import os
+    import random
 
-    if args.multi>1:
+    if args.multi > 1:
         setup_multi(args)
         return 0
 
-    num_jobs=0
+    num_jobs = 0
     for folder in os.listdir(args.calc_dir):
-        if folder[0:15]=="Local_TraceWin_":
-            num_jobs+=1
+        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)
+    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 range(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
+    if args.num_jobs > num_jobs:
+        for n in range(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
+    stat_study = args.study + 13
 
     create_and_copy_lattice(args, args.calc_dir)
 
     if not os.path.isdir("TraceWin/exe"):
         setup_exec_folder()
 
-    if args.seed!=None:
+    if args.seed is not None:
         random.seed(args.seed)
-    rand_lst=[random.randint(1e5,1e9) for i in range(num_jobs)]
-    #rand_lst=file('rand.dat','r').read().split()
+    rand_lst = [random.randint(1e5, 1e9) for i in range(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"))
+    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)
+        if input_files[i][-2:] == "_0":
+            input_files[i] = input_files[i][:-2] + "_$(Process)"
+    input_files = ", ".join(input_files)
 
-    head_template=__queue_templates__.get_head_template()
-    queue_template=__queue_templates__.get_queue_template()
+    head_template = __queue_templates__.get_head_template()
+    queue_template = __queue_templates__.get_queue_template()
     with open(args.calc_dir + "/submit.job", "w") as out:
         out.write(head_template.format(**locals()))
 
         for i in range(num_jobs):
-            q_priority=max(-5,args.priority-i/100)
-            rand=rand_lst[i]
+            q_priority = max(-5, args.priority - i / 100)
+            rand = rand_lst[i]
             out.write(queue_template.format(**locals()))
 
+
 def ask(args):
-    print("You have now set up an error study running {} jobs in folder {}".format(args.num_jobs,args.calc_dir))
-    if args.multi>1:
+    print("You have now set up an error study running {} jobs in folder {}".format(args.num_jobs, args.calc_dir))
+    if args.multi > 1:
         print("This is a multi-error study")
-    answer  = input("Do you want to submit the jobs immediately? (yes/no) ")
-    if answer=='yes':
+    answer = input("Do you want to submit the jobs immediately? (yes/no) ")
+    if answer == "yes":
         return 0
 
     import sys
+
     sys.exit(0)
-    
-def execute(calc_dir,multi):
+
+
+def execute(calc_dir, multi):
     import subprocess
+
     if multi:
-        cmd=["condor_submit",calc_dir+"/multi.job"]
+        cmd = ["condor_submit", calc_dir + "/multi.job"]
     else:
-        cmd=["condor_submit",calc_dir+"/submit.job"]
+        cmd = ["condor_submit", calc_dir + "/submit.job"]
     subprocess.call(cmd)
 
-if __name__=="__main__":
-    args=define_and_handle_args()
+
+if __name__ == "__main__":
+    args = define_and_handle_args()
     setup(args)
     ask(args)
-    execute(args.calc_dir,args.multi>1)
-
+    execute(args.calc_dir, args.multi > 1)