From e419f93b6931e96d7976d69b0718b49e4fe0087f Mon Sep 17 00:00:00 2001
From: Yngve Levinsen <yngve.levinsen@esss.se>
Date: Tue, 28 Aug 2018 13:52:29 +0200
Subject: [PATCH] fix for when lattice file is in different subfolder and a
 little import cleaning

---
 ess/lib_tw.py | 23 +++++++++++++----------
 1 file changed, 13 insertions(+), 10 deletions(-)

diff --git a/ess/lib_tw.py b/ess/lib_tw.py
index 7d280d6..660d2ad 100644
--- a/ess/lib_tw.py
+++ b/ess/lib_tw.py
@@ -18,9 +18,7 @@ from struct    import pack
 from itertools import chain
 
 from subprocess import check_output
-from os         import system
-from os.path    import isdir,isfile
-from sys        import exit
+import os,sys
 
 from .lib_tw_elem import *
 
@@ -37,7 +35,7 @@ class LATTICE:
         '''
         :param file_name_lat: name of lattice file
         :type file_name_lat: str
-        :param file_name_fmap: list of file map file(-s)
+        :param file_name_fmap: list of field map file(-s)
         :type file_name_fmap: list or str
         :param freq: RF frequency
         :type freq: float
@@ -46,7 +44,8 @@ class LATTICE:
         '''
 
         # In case file_name_fmap is str
-        if isinstance(file_name_fmap,str): file_name_fmap=[file_name_fmap]
+        if isinstance(file_name_fmap,str):
+            file_name_fmap=[file_name_fmap]
 
         # Elem/comm class dict
         dic_cls={'DRIFT'               : DRIFT               ,
@@ -81,8 +80,12 @@ class LATTICE:
         # Field map dict
         dic_fmap={}
         for file_name_fmap_i in file_name_fmap:
-            name_fmap='.'.join(file_name_fmap_i.split('/')[-1].split('.')[:-1])  # Remove / and extention
+            name_fmap='.'.join(file_name_fmap_i.split('/')[-1].split('.')[:-1])  # Remove / and extension
             dic_fmap[name_fmap]=FIELD_MAP_DATA(file_name_fmap_i)
+        # In case the lattice file is in a different folder:
+        basefolder=os.path.dirname(file_name_lat)
+        if basefolder:
+            _update_field_map_dict( dic_fmap, basefolder)
 
         # Go through the lat file
         with open(file_name_lat) as file:
@@ -361,9 +364,9 @@ class PROJECT:
         if self.flag_hide    !=None: opt_exe+=' hide'
 
         ## if self.path_cal!=None:
-        ##     if isdir(self.path_cal)==False: system('mkdir '+self.path_cal)
+        ##     if os.isdir(self.path_cal)==False: system('mkdir '+self.path_cal)
 
-        system(opt_exe)
+        os.system(opt_exe)
 
 #---- Data related
 
@@ -889,7 +892,7 @@ def loss_elem2den(s,loss,file_name_dt='',dlt_dt=5e-6):
 
     return loss_den
 
-    exit()
+    sys.exit()
 
     # Define L
     l=[s[i+1]-s[i] for i in range(len(s)-1)]; l.insert(0,0.0)
@@ -917,7 +920,7 @@ def loss_elem2den(s,loss,file_name_dt='',dlt_dt=5e-6):
         if l[i]==0.0 and loss[i]!=0.0:
             if i==1:
                 print('The 1st elem has 0 length and yet has a loss, exiting...')
-                exit()
+                sys.exit()
             else:
                 k=1
                 while l[i-k]==0.0: k+=1
-- 
GitLab