diff --git a/ess/lib_tw.py b/ess/lib_tw.py index 5fa7664ffd6723f9bad2bddb4ba3cf4038f836e9..f19eb63c9261302b34d881548be28e1a238d3a81 100644 --- a/ess/lib_tw.py +++ b/ess/lib_tw.py @@ -344,6 +344,21 @@ class LATTICE: except AttributeError: pass + def get_bdsim(self,output_folder='bdsim'): + if not os.path.exists(output_folder): + os.makedirs(output_folder) + file_name=os.path.join(output_folder,'bdsim.dat') + + if self.lst[-1].gamma==1.0: self.update_gamma() # Assign gamma, if not done yet + + with open(file_name,'w') as fname: + for fmap in self.fmap: + fname.write(self.fmap[fmap].get_bdsim(os.path.join(output_folder,fmap))+'\n') + for lat_i in self.lst: + if hasattr(lat_i,'get_bdsim'): + fname.write(lat_i.get_bdsim()+'\n') + + class PROJECT: ''' - This is for running multiple simulations 1-by-1 under 1 project. diff --git a/ess/lib_tw_elem.py b/ess/lib_tw_elem.py index 7ff762c82333e53215ac4a2c5ce467a79ebaa7db..05ac99a954da0fa71195042ab4e353595dc2a0a7 100644 --- a/ess/lib_tw_elem.py +++ b/ess/lib_tw_elem.py @@ -98,6 +98,7 @@ from __future__ import print_function from numpy import * import sys +import shutil #---- Constants @@ -122,7 +123,15 @@ class FIELD_MAP_DATA: # Instances self.dL =float(lin[0].split()[1])/int(lin[0].split()[0]) - self.field=map(float,lin[2:]) + self.field=array(lin[2:], dtype=float) + self.file_name = file_name + + def get_bdsim(self,path): + # Only support 1D edz for now! + with open(path+'.txt', 'w') as fmapout: + for i in range(len(self.field)): + fmapout.write('{} {} {} {}\n'.format(self.dL*i, 0.0, 0.0, self.field[i])) + return '{0}: field, type="ebmap2d", integrator=g4classicalrk4, electricFile=bdsim1d:{0}.txt;'.format(self.file_name[:-4]) #---- Parent classes @@ -298,6 +307,9 @@ class FIELD_MAP(ELEM): lin+='0 0 0 0 0' return lin + def get_bdsim(self): + + return 'element, geometryFile="gdml:{0}.gdml", fieldVacuum="{0}", l={1}'.format(self.name_fmap,self.L) #---- Active elements @@ -361,6 +373,10 @@ class DRIFT(ELEM): return lin + def get_bdsim(self): + + return 'drift, l={}*mm, beampipeRadius={}*mm;'.format(self.L,self.apt) + class QUAD(ELEM): ''' - Tilt not supported. @@ -433,6 +449,11 @@ class QUAD(ELEM): return lin + + def get_bdsim(self): + + return 'quadrupole, l={}*mm, k1={}, beampipeRadius={}*mm;'.format(self.L, self.G/Brho(self.gamma), self.apt) + class THIN_STEERING(ELEM): ''' - Only magnetic steerer supported.