diff --git a/ess/lib_tw.py b/ess/lib_tw.py index 96d176ef99c2cdf57608ab6f789337dfa4b55d2a..702a95b3d8addb0e023fc5d5973f8f30f99c19c8 100644 --- a/ess/lib_tw.py +++ b/ess/lib_tw.py @@ -33,7 +33,7 @@ class LATTICE: ''' 2015.10.14 ''' - def __init__(self,file_name_lat,file_name_fmap,freq=352.21,gamma=1.0): + def __init__(self,file_name_lat,file_name_fmap=[],freq=352.21,gamma=1.0): # In case file_name_fmap is str if isinstance(file_name_fmap,basestring): file_name_fmap=[file_name_fmap] @@ -68,10 +68,11 @@ class LATTICE: typ =lin.split()[0 ].upper() para=lin.split()[1:] # Map to a class - if typ == 'FIELD_MAP' : lst.append( FIELD_MAP(name,typ,para,dic_fmap)) - elif typ in dic_cls.keys(): lst.append(dic_cls[typ](name,typ,para) ) - elif 'DIAG' in typ : lst.append( DIAG(name,typ,para) ) - else : lst.append( COMM(name,typ,para) ) + if typ == 'FIELD_MAP' : lst.append( FIELD_MAP(name,typ,para,dic_fmap)) + elif typ == 'FIELD_MAP_PATH' : _update_field_map_dict(dic_fmap,para[0]) + elif typ in dic_cls.keys() : lst.append(dic_cls[typ](name,typ,para) ) + elif 'DIAG' in typ : lst.append( DIAG(name,typ,para) ) + else : lst.append( COMM(name,typ,para) ) # Break the loop if typ=='END': break @@ -938,6 +939,17 @@ def partran_end_all(file_name_in_all,file_name_out): else : file_out.write(' %10s'%lin[i]) file_out.write('\n') + + +def _update_field_map_dict(dictionary,folder_path): + import os + for filename in os.listdir(folder_path): + if filename.split('.')[-1]=='edz': # only 1D for now.. + key=filename.split('.')[0] + if key not in dictionary: # do not override user selection + dictionary[key]=os.path.join(folder_path,filename) + print key,':',dictionary[key] + #-------- Obsolete classes and functions #---- Old MATCH and LATTICE classes from the time of IPAC'15 (more complex) diff --git a/ess/lib_tw_elem.py b/ess/lib_tw_elem.py index 1f58ee969fea5833ed593f260819520b230cc1d9..a9da63b786d7805b4adaaadbf85a3ef843b38c3b 100644 --- a/ess/lib_tw_elem.py +++ b/ess/lib_tw_elem.py @@ -90,6 +90,9 @@ class FIELD_MAP_DATA: ''' def __init__(self,file_name): + if file_name.split('.')[-1]!='edz': + raise ValueError('Only 1D field maps (Ez) is supported for now') + # Read file with open(file_name) as file: lin=file.readlines()