diff --git a/ess/lib_tw.py b/ess/lib_tw.py
new file mode 100644
index 0000000000000000000000000000000000000000..dd0a8800da79bbbf12df59e410940a225cd60800
--- /dev/null
+++ b/ess/lib_tw.py
@@ -0,0 +1,1129 @@
+
+'''
+ Classes and functions for TraceWin
+
+ - Note for clean-up
+
+ * Old LATTICE and MATCH classes can be deleted once trajectory
+ correction is established with new classes.
+'''
+
+#---- Lib
+
+from numpy import *
+from numpy.linalg import det
+
+from struct import pack
+from itertools import chain
+
+from subprocess import check_output
+from os import system
+from os.path import isdir,isfile
+
+from lib_tw_elem import *
+
+
+
+#-------- Classes
+
+#---- Lattice and project
+
+class LATTICE:
+ '''
+ 2015.10.14
+ '''
+ 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)==True: file_name_fmap=[file_name_fmap]
+
+ # Elem/comm class dict
+ dic_cls={'DRIFT':DRIFT,
+ 'QUAD':QUAD,'BEND':BEND,'EDGE':EDGE,'THIN_STEERING':THIN_STEERING,
+ 'GAP':GAP,'DTL_CEL':DTL_CEL,
+ 'APERTURE':APERTURE,'DIAG_POSITION':DIAG_POSITION,
+ 'STEERER':STEERER,'CHOPPER':CHOPPER,
+ 'FREQ':FREQ,'MARKER':MARKER}
+
+ # 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
+ dic_fmap[name_fmap]=FIELD_MAP_DATA(file_name_fmap_i)
+
+ # Go through the lat file
+ with open(file_name_lat) as file:
+ lst=[]
+ for lin in file:
+ lin=lin.partition(';')[0] # Remove comment
+ if lin.split()!=[]: # Remove empty line
+ # Split a line
+ if ':' in lin:
+ name=lin.partition(':')[0].split()[0 ]
+ typ =lin.partition(':')[2].split()[0 ].upper()
+ para=lin.partition(':')[2].split()[1:]
+ else:
+ name=''
+ 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) )
+ # Break the loop
+ if typ=='END': break
+
+ # Instances
+ self.gamma=gamma
+ self.freq =freq
+ self.lst =lst
+ self.fmap =dic_fmap # Maybe not needed ...
+
+ # Assign idx, idx_elem, s, freq, apt
+ self.update_idx()
+
+ def update_idx(self):
+
+ # Assign/update idx, idx_elem, s, freq
+ for i in range(len(self.lst)):
+ if i==0:
+ self.lst[0].freq=self.freq
+ if i!=0:
+ self.lst[i].idx =self.lst[i-1].idx
+ self.lst[i].idx_elem=self.lst[i-1].idx_elem
+ self.lst[i].s =self.lst[i-1].s
+ self.lst[i].freq =self.lst[i-1].freq
+ self.lst[i].update_idx()
+ # Assign apt (using apt of the previous elem for diag elem)
+ for i in range(len(self.lst)):
+ try:
+ if self.lst[i].apt==None:
+ if self.lst[i].idx_elem==0:
+ for k in range(1,len(self.lst)):
+ try:
+ if self.lst[k].apt!=None:
+ self.lst[i].apt=self.lst[k].apt; break
+ except: pass
+ else:
+ for k in range(i)[::-1]:
+ try : self.lst[i].apt=self.lst[k].apt; break
+ except: pass
+ except: pass
+
+ def update_gamma(self):
+
+ for i in range(len(self.lst)):
+ if i==0: self.lst[0].gamma=self.gamma
+ if i!=0: self.lst[i].gamma=self.lst[i-1].gamma
+ try : self.lst[i].update_gamma()
+ except: pass
+
+ def update_st(self,file_name):
+ '''
+ Assign/update steerer values from "Steerer_Values.txt".
+ '''
+ # Extract BLx and BLy from the file
+ with open(file_name,'r') as file:
+ BLx={}; BLy={}
+ for lin in file:
+ lin=lin.split()[3:]
+ for i in range(len(lin)):
+ if lin[i]==':' : idx_elem =int(lin[i-1])-1 # "-1" since idx=0,1,2,...
+ if lin[i]=='BLx=': BLx[idx_elem]=float(lin[i+1])
+ if lin[i]=='BLy=': BLy[idx_elem]=float(lin[i+1])
+
+ # Assign/update steerers
+ for i in range(len(self.lst)):
+ if self.lst[i].idx_elem in BLx.keys():
+ idx_elem=self.lst[i].idx_elem
+ if self.lst[i].typ=='THIN_STEERING':
+ self.lst[i].BLx=BLx[idx_elem]
+ self.lst[i].BLy=BLy[idx_elem]
+ if self.lst[i].typ!='THIN_STEERING':
+ for k in range(i)[::-1]:
+ if self.lst[k].typ=='STEERER':
+ self.lst[k].Bx=BLx[idx_elem]/self.lst[i].L
+ self.lst[k].By=BLy[idx_elem]/self.lst[i].L
+ break
+
+ def get_tw(self,file_name):
+
+ with open(file_name,'w') as file:
+ for lat_i in self.lst:
+ file.write(lat_i.get_tw()+'\n')
+
+ def get_madx(self,file_name_elem='elem.madx',file_name_seq='seq.madx'):
+
+ if self.lst[-1].gamma==1.0: self.update_gamma() # Assign gamma, if not done yet
+
+ with open(file_name_elem,'w') as file:
+ for lat_i in self.lst:
+ try : print >>file,lat_i.get_madx()
+ except: pass
+
+ with open(file_name_elem,'r') as file: lst_name=[lin.split(':')[0] for lin in file]
+ with open(file_name_seq ,'w') as file: print >>file,'linac:line=('+','.join(lst_name)+');'
+
+ def get_fluka(self,file_name='elem.dat'):
+
+ if self.lst[-1].gamma==1.0: self.update_gamma() # Assign gamma, if not done yet
+
+ with open(file_name,'w') as file:
+ for lat_i in self.lst:
+ try : print >>file,lat_i.get_fluka()
+ except: pass
+
+ def get_mars(self,file_name='elem.dat'):
+
+ if self.lst[-1].gamma==1.0: self.update_gamma() # Assign gamma, if not done yet
+
+ with open(file_name,'w') as file:
+ for lat_i in self.lst:
+ try : print >>file,lat_i.get_mars()
+ except: pass
+
+class PROJECT:
+ '''
+ - This is for running multiple simulations 1-by-1 under 1 project.
+ - Maybe not very useful...
+
+ 2015.10.15
+ '''
+
+ def __init__(self,file_name_proj='proj.ini'):
+
+ # Instances (Add more options as needed.)
+ self.file_name_proj=file_name_proj
+ self.file_name_lat =None
+ self.path_cal =None
+ self.seed =None
+ self.flag_hide =None
+
+ def exe(self):
+
+ opt_exe ='TraceWin64_noX11 '+self.file_name_proj
+ if self.file_name_lat!=None: opt_exe+=' dat_file=' +self.file_name_lat
+ if self.path_cal !=None: opt_exe+=' path_cal=' +self.path_cal
+ if self.seed !=None: opt_exe+=' random_seed='+self.seed
+ if self.flag_hide !=None: opt_exe+=' hide'
+
+ ## if self.path_cal!=None:
+ ## if isdir(self.path_cal)==False: system('mkdir '+self.path_cal)
+
+ system(opt_exe)
+
+#---- Data related
+
+class PARTRAN:
+ '''
+ - The list not complete. Add parameters as needed.
+
+ 2015.11.24
+ '''
+ def __init__(self,file_name):
+
+ # Consts to convert phs to z.
+ mass=938.272029
+ c =2.99792458
+ freq=352.21
+
+ # Extract indices.
+ with open(file_name) as file:
+ for lin in file.readlines():
+ lin=lin.split()
+ if lin[0]=='####':
+ idx_s =lin.index("z(m)" ); idx_gamma=lin.index("gama-1")
+ idx_x =lin.index("x0" ); idx_y =lin.index("y0" ); idx_phs =lin.index("p0" )
+ idx_sigx=lin.index("SizeX" ); idx_sigy =lin.index("SizeY" ); idx_sigz=lin.index("SizeZ"); idx_sigp=lin.index("SizeP")
+ idx_alfx=lin.index("sxx'" ); idx_alfy =lin.index("syy'" ); idx_alfz=lin.index("szdp" )
+ idx_epsx=lin.index("ex" ); idx_epsy =lin.index("ey" ); idx_epsz=lin.index("ezdp" ); idx_epsp=lin.index("ep" )
+ idx_halx=lin.index("hx" ); idx_haly =lin.index("hy" ); idx_halp=lin.index("hp" )
+ idx_loss=lin.index("Powlost")
+ break
+
+ # Extract data.
+ with open(file_name) as file:
+ data=[]; flag=0
+ for lin in file.readlines():
+ lin=lin.split()
+ if flag ==1 : data.append(map(float,lin))
+ if lin[0]=='####': flag=1
+ data=array(data).transpose()
+
+ # Instances
+ self.s =data[idx_s ]
+ self.x =data[idx_x ]; self.y =data[idx_y ]; self.phs =data[idx_phs ]
+ self.sigx=data[idx_sigx]; self.sigy=data[idx_sigy]; self.sigz=data[idx_sigz]; self.sigp=data[idx_sigp]
+ self.epsx=data[idx_epsx]; self.epsy=data[idx_epsy]; self.epsz=data[idx_epsz]; self.epsp=data[idx_epsp]
+ self.halx=data[idx_halx]; self.haly=data[idx_haly]; self.halz=data[idx_halp]; self.halp=data[idx_halp]
+ self.loss=data[idx_loss]
+
+ # Additional instances
+ self.gamma= data[idx_gamma]+1.0
+ self.beta = sqrt(1.0-1.0/self.gamma**2)
+ self.z =-self.phs*self.beta*(c/freq*1e5)/360.0
+ self.betx = self.sigx**2/self.epsx*self.beta*self.gamma
+ self.bety = self.sigy**2/self.epsy*self.beta*self.gamma
+ self.betz = self.sigz**2/self.epsz*self.beta*self.gamma**3
+ self.betp = self.sigp**2/self.epsp
+ self.alfx =-data[idx_alfx]/self.epsx*self.beta*self.gamma
+ self.alfy =-data[idx_alfy]/self.epsy*self.beta*self.gamma
+ self.alfz =-data[idx_alfz]/self.epsz*self.beta*self.gamma**3
+ self.alfp =-self.alfz
+
+ # Convert to list (not necessary?)
+ self.s =self.s.tolist() ; self.gamma=self.gamma.tolist(); self.beta=self.beta.tolist()
+ self.x =self.x.tolist() ; self.y =self.y.tolist() ; self.z =self.z.tolist() ; self.phs =self.phs.tolist()
+ self.sigx=self.sigx.tolist(); self.sigy =self.sigy.tolist() ; self.sigz=self.sigz.tolist(); self.sigp=self.sigp.tolist()
+ self.betx=self.betx.tolist(); self.bety =self.bety.tolist() ; self.betz=self.betz.tolist(); self.betp=self.betp.tolist()
+ self.alfx=self.alfx.tolist(); self.alfy =self.alfy.tolist() ; self.alfz=self.alfz.tolist(); self.alfp=self.alfp.tolist()
+ self.epsx=self.epsx.tolist(); self.epsy =self.epsy.tolist() ; self.epsz=self.epsz.tolist(); self.epsp=self.epsp.tolist()
+ self.halx=self.halx.tolist(); self.haly =self.haly.tolist() ; self.halz=self.halz.tolist(); self.halp=self.halp.tolist()
+ self.loss=self.loss.tolist()
+
+ def loss_den(self,file_name_dt,dlt_dt=5e-6):
+
+ return loss_elem2den(self.s,self.loss,file_name_dt,dlt_dt)
+
+class DST:
+ '''
+ Class for a TraceWin's .dst file.
+
+ - TraceWin seems using beta and gamma for each particle
+ so the conversion to (z,z') is based on this assumption.
+
+ 2015.10.06
+ '''
+ def __init__(self,file_name,unit_x='cm',unit_px='rad',unit_z='rad',unit_pz='MeV'):
+
+ # Const
+ c=2.99792458
+
+ # Read the file
+ with open(file_name) as file:
+ dummy=fromfile(file,dtype= uint8 ,count=2 )
+ Nptcl=fromfile(file,dtype= uint32,count=1 )[0]
+ Ibeam=fromfile(file,dtype=float64,count=1 )[0]
+ freq =fromfile(file,dtype=float64,count=1 )[0]
+ dummy=fromfile(file,dtype= uint8 ,count=1 )
+ x =fromfile(file,dtype=float64,count=Nptcl*6).reshape(Nptcl,6).transpose()
+ mass =fromfile(file,dtype=float64,count=1 )[0]
+
+ # Adjust units
+ gamma=1.0+x[5]/mass; beta=sqrt(1-1/gamma**2)
+ if unit_x =='mm' : x[0]= x[0]*1e1; x[2]=x[2]*1e1
+ if unit_px=='mrad': x[1]= x[1]*1e3; x[3]=x[3]*1e3
+ if unit_z =='deg' : x[4]= x[4]*180/pi
+ if unit_z =='mm' : x[4]=-x[4]*c*beta/(2*pi*freq)*1e5
+ if unit_pz=='mrad': x[5]=(x[5]-mean(x[5]))/(mass*beta**2*gamma**3)*1e3
+
+ # Instances
+ self.x =x.transpose()
+ self.mass =mass
+ self.freq =freq
+ self.Ibeam=Ibeam
+
+class DENSITY:
+ '''
+ - Note instances are not identical for Nrun=1 and Nrun>1.
+ - Be careful with # of steps for an envelope file.
+ - When Nrun>1, ave and rms in a file are sum and squared sum.
+ - Double check before a production !!!!
+ - Dim of arrays:
+
+ Nelem(Nidx): idx_elem, s, Nptcl, Ibeam
+
+ 2 x Nelem(Nidx): apt # x, y
+ 3 x Nelem(Nidx): accpt # phs+, phs-, ken
+ 7 x Nelem(Nidx): cent_(..), sig, xmax, xmin # x, y, phs, ken, r, z, dpp
+ 3 x Nelem(Nidx): eps # x, y, phs
+
+ Nrun x Nelem: loss
+ Nelem : loss_num_(..), loss_pow_(..)
+
+ Nelem x Nstep: den
+
+ 2015.11.12
+ '''
+ def __init__(self,file_name):
+
+ #-- Empty arrays
+
+ idx_elem =[]; s =[]
+ apt =[]; accpt =[]
+ Nptcl =[]; Ibeam =[]
+ cent_ave =[]; cent_rms =[]
+ cent_max =[]; cent_min =[]
+ sig_ave =[]; sig_rms =[]
+ eps_ave =[]; eps_rms =[]
+ xmax =[]; xmin =[]
+ loss_num =[]; loss_pow =[]
+ loss_num_ave=[]; loss_pow_ave=[]
+ loss_num_rms=[]; loss_pow_rms=[]
+ loss_num_max=[]; loss_pow_max=[]
+ loss_num_min=[]; loss_pow_min=[]
+ den =[]; den_pow =[]
+
+ #-- Extract data
+
+ with open(file_name) as file:
+ while True:
+ try:
+ # Partran and envelope
+ ver,year,flag_long=fromfile(file,dtype=uint16,count=3)
+ Nrun =fromfile(file,dtype=uint32,count=1)[0]
+ idx_elem.append(fromfile(file,dtype= uint32,count=1)[0])
+ Ibeam.append(fromfile(file,dtype=float32,count=1)[0])
+ s.append( fromfile(file,dtype=float32,count=1)[0])
+ apt.append( fromfile(file,dtype=float32,count=2) )
+ Nstep=fromfile(file,dtype=uint32,count=1)[0]
+ cent_ave.append(fromfile(file,dtype=float32,count=7))
+ cent_rms.append(fromfile(file,dtype=float32,count=7))
+ xmax.append( fromfile(file,dtype=float32,count=7))
+ xmin.append( fromfile(file,dtype=float32,count=7))
+ if ver> 5:
+ sig_ave.append(fromfile(file,dtype=float32,count=7))
+ sig_rms.append(fromfile(file,dtype=float32,count=7))
+ if ver>=6:
+ cent_min.append(fromfile(file,dtype=float32,count=7))
+ cent_max.append(fromfile(file,dtype=float32,count=7))
+ if ver>=7:
+ eps_ave.append(fromfile(file,dtype=float32,count=3))
+ eps_rms.append(fromfile(file,dtype=float32,count=3))
+ if ver>=8:
+ accpt.append(fromfile(file,dtype=float32,count=3))
+ Nptcl.append(fromfile(file,dtype=uint64,count=1)[0])
+ # Partran only
+ if Nptcl[-1]>0:
+ loss_num.append([]); loss_pow.append([]); den.append([]); den_pow.append([])
+ for n in range(Nrun):
+ loss_num[-1].append(fromfile(file,dtype= uint64,count=1)[0])
+ loss_pow[-1].append(fromfile(file,dtype=float32,count=1)[0])
+ loss_num_ave.append(sum(loss_num[-1]))
+ loss_num_rms.append(fromfile(file,dtype= uint64,count=1)[0])
+ loss_num_min.append(fromfile(file,dtype= uint64,count=1)[0])
+ loss_num_max.append(fromfile(file,dtype= uint64,count=1)[0])
+ loss_pow_ave.append(sum(loss_pow[-1]))
+ loss_pow_rms.append(fromfile(file,dtype=float64,count=1)[0])
+ loss_pow_min.append(fromfile(file,dtype=float32,count=1)[0])
+ loss_pow_max.append(fromfile(file,dtype=float32,count=1)[0])
+ for k in range(7):
+ if flag_long==1: den[-1].append(fromfile(file,dtype=uint64,count=Nstep))
+ else : den[-1].append(fromfile(file,dtype=uint32,count=Nstep))
+ if Ibeam[-1]>0:
+ for k in range(3):
+ den_pow[-1].append(fromfile(file,dtype=float32,count=Nstep))
+ #print Nrun,Nptcl[-1],idx_elem[-1] # Diag
+ except:
+ break
+
+ #-- Reshape arrays
+
+ apt =swapaxes( apt,1,0); accpt =swapaxes( accpt,1,0)
+ cent_ave=swapaxes(cent_ave,1,0); cent_rms=swapaxes(cent_rms,1,0)
+ cent_max=swapaxes(cent_max,1,0); cent_min=swapaxes(cent_min,1,0)
+ sig_ave =swapaxes( sig_ave,1,0); sig_rms =swapaxes( sig_rms,1,0)
+ eps_ave =swapaxes( eps_ave,1,0); eps_rms =swapaxes( eps_rms,1,0)
+ xmax =swapaxes( xmax,1,0); xmin =swapaxes( xmin,1,0)
+ if Nptcl[0]>0:
+ loss_num=swapaxes(loss_num,1,0); loss_pow=swapaxes(loss_pow,1,0)
+ den =swapaxes( den,1,0); den_pow =swapaxes( den_pow,1,0)
+
+ #-- Take care ave and rms
+
+ cent_ave=cent_ave/Nrun; cent_rms=sqrt(cent_rms/Nrun)
+ sig_ave = sig_ave/Nrun; sig_rms =sqrt( sig_rms/Nrun)
+ eps_ave = eps_ave/Nrun; eps_rms =sqrt( eps_rms/Nrun)
+ if Nptcl[0]>0:
+ loss_num_ave=1.0*array(loss_num_ave)/Nrun; loss_num_rms=sqrt(1.0*array(loss_num_rms)/Nrun)
+ loss_pow_ave= array(loss_pow_ave)/Nrun; loss_pow_rms=sqrt( array(loss_pow_rms)/Nrun)
+
+ #-- Change units, m => mm, pi-m-rad => pi-mm-mrad
+
+ apt *=1e3
+ eps_ave*=1e6; eps_rms*=1e6
+ for k in (0,1,4,5):
+ cent_ave[k]*=1e3; cent_rms[k]*=1e3
+ cent_max[k]*=1e3; cent_min[k]*=1e3
+ sig_ave[k] *=1e3; sig_rms[k] *=1e3
+ xmax[k] *=1e3; xmin[k] *=1e3
+
+ #-- Define std (around to avoid sqrt(-eps))
+
+ if Nrun>1:
+ cent_std=sqrt(around(cent_rms**2-cent_ave**2,12))
+ sig_std =sqrt(around( sig_rms**2- sig_ave**2,12))
+ eps_std =sqrt(around( eps_rms**2- eps_ave**2,12))
+ cent_std=nan_to_num(cent_std) # Replace nan with 0
+ sig_std =nan_to_num( sig_std) # Replace nan with 0
+ eps_std =nan_to_num( eps_std) # Replace nan with 0
+ if Nptcl[0]>0:
+ loss_num_std=sqrt(around(loss_num_rms**2-loss_num_ave**2,16))
+ loss_pow_std=sqrt(around(loss_pow_rms**2-loss_pow_ave**2,16))
+ loss_num_std=nan_to_num(loss_num_std) # Replace nan with 0
+ loss_pow_std=nan_to_num(loss_pow_std) # Replace nan with 0
+
+ #-- Convert to list (just in case...)
+
+ apt = apt.tolist(); accpt = accpt.tolist(); accpt.append(accpt[0]); del accpt[0]
+ cent_ave=cent_ave.tolist(); cent_rms=cent_rms.tolist()
+ cent_max=cent_max.tolist(); cent_min=cent_min.tolist()
+ sig_ave = sig_ave.tolist(); sig_rms = sig_rms.tolist()
+ eps_ave = eps_ave.tolist(); eps_rms = eps_rms.tolist()
+ xmax = xmax.tolist(); xmin=xmin.tolist()
+ if Nptcl[0]>0:
+ loss_num = loss_num.tolist(); loss_pow = loss_pow.tolist()
+ loss_num_ave=loss_num_ave.tolist(); loss_pow_ave=loss_pow_ave.tolist()
+ loss_num_rms=loss_num_rms.tolist(); loss_pow_rms=loss_pow_rms.tolist()
+ den = den.tolist(); den_pow = den_pow.tolist()
+
+ if Nrun>1:
+ cent_std=cent_std.tolist()
+ sig_std = sig_std.tolist()
+ eps_std = eps_std.tolist()
+ if Nptcl[0]>0:
+ loss_num_std=loss_num_std.tolist()
+ loss_pow_std=loss_pow_std.tolist()
+
+ #-- Outputs
+
+ self.idx_elem=idx_elem
+ self.s =s
+ self.apt =apt
+ self.accpt =accpt
+ self.Nptcl =Nptcl
+ self.Ibeam =Ibeam
+
+ self.Nrun =Nrun
+ self.Nstep=Nstep
+
+ if Nrun==1:
+ self.cent=cent_ave
+ self.sig = sig_ave
+ self.eps = eps_ave
+ self.xmax= xmax
+ self.xmin= xmin
+ if Nptcl[0]>0:
+ self.loss_num=loss_num[0]; self.loss_pow=loss_pow[0]
+ self.den = den; self.den_pow = den_pow
+ else:
+ self.cent_ave=cent_ave; self.cent_rms=cent_rms; self.cent_std=cent_std
+ self.cent_max=cent_max; self.cent_min=cent_min
+ self.sig_ave = sig_ave; self.sig_rms = sig_rms; self.sig_std = sig_std
+ self.eps_ave = eps_ave; self.eps_rms = eps_rms; self.eps_std = eps_std
+ self.xmax = xmax; self.xmin = xmin
+ if Nptcl[0]>0:
+ self.loss_num =loss_num ; self.loss_pow =loss_pow
+ self.loss_num_ave=loss_num_ave; self.loss_pow_ave=loss_pow_ave
+ self.loss_num_rms=loss_num_rms; self.loss_pow_rms=loss_pow_rms
+ self.loss_num_std=loss_num_std; self.loss_pow_std=loss_pow_std
+ self.loss_num_max=loss_num_max; self.loss_pow_max=loss_pow_max
+ self.loss_num_min=loss_num_min; self.loss_pow_min=loss_pow_min
+ self.den = den ; self.den_pow = den_pow
+
+ #-- Option outputs
+
+ self.idx_4_elem_end=[len(idx_elem)-1-idx_elem[::-1].index(i) for i in list(set(idx_elem))]
+
+#-------- Functions
+
+#---- Dist related
+
+def x2dst(x,mass,freq,Ibeam,path_file='part_dtl1_new.dst'):
+ '''
+ Output a TraceWin's .dst file from x and etc.
+
+ Input: x (Nptcl,6)
+
+ For binary characters see https://docs.python.org/2/library/struct.html
+
+ 2014.10.03
+ '''
+
+ file=open(path_file,'w')
+ out =pack('b',125)
+ out+=pack('b',100)
+ out+=pack('i',len(x))
+ out+=pack('d',Ibeam)
+ out+=pack('d',freq)
+ out+=pack('b',125) # Typo in the manual !!!!
+ x=list(chain(*x)) # Flatten x
+ for x_i in x: out+=pack('d',x_i)
+ out+=pack('d',mass)
+ print >>file,out
+ file.close()
+
+def plt2x(path_file):
+ '''
+ Extract x and etc from a TraceWin's binary .plt file.
+ The units are (cm,rad,cm,rad,rad,MeV,loss).
+
+ Output: x (Nelem,Nptcl,7)
+
+ For binary characters see https://docs.python.org/2/library/struct.html
+
+ 2014.10.06
+ '''
+
+ file=open(path_file,'r')
+
+ # Hetero info
+ file.read(1); file.read(1) # 125,100
+ Nelem=unpack('i',file.read(4))[0]
+ Nptcl=unpack('i',file.read(4))[0]
+ Ibeam=unpack('d',file.read(8))[0]
+ freq =unpack('d',file.read(8))[0]
+ mass =unpack('d',file.read(8))[0]
+
+ # Loop for elem
+ x=[]; i_unk=[]; i_elem=[]; s=[]; phs=[]; ken=[];
+ for i in range(Nelem):
+ try:
+ i_unk.append(unpack('b',file.read(1))[0])
+ i_elem.append(unpack('i',file.read(4))[0])
+ s.append(unpack('d',file.read(8))[0])
+ phs.append(unpack('d',file.read(8))[0])
+ ken.append(unpack('d',file.read(8))[0])
+ x.append([[unpack('f',file.read(4))[0] for l in range(7)] for k in range(Nptcl)])
+ except:
+ break
+
+ file.close()
+
+ return x,mass,freq,Ibeam,i_unk,i_elem,s,phs,ken
+
+def x2plt(x,mass,freq,Ibeam,i_unk,i_elem,s,phs,ken,path_file='dtl1_new.plt'):
+ '''
+ Output a TraceWin's .plt file from x and etc.
+
+ Input: x (Nelem,Nptcl,7)
+
+ For binary characters see https://docs.python.org/2/library/struct.html
+
+ 2014.10.07
+ '''
+
+ file=open(path_file,'w')
+ out =pack('b',125)
+ out+=pack('b',100)
+ out+=pack('i',len(x))
+ out+=pack('i',len(x[0]))
+ out+=pack('d',Ibeam)
+ out+=pack('d',freq)
+ out+=pack('d',mass)
+ for i in range(len(x)):
+ out+=pack('b',i_unk[i])
+ #out+=pack('i',i_elem[i])
+ out+=pack('i',i) # To truncate the elem number
+ out+=pack('d',s[i])
+ out+=pack('d',phs[i])
+ out+=pack('d',ken[i])
+ x_i=list(chain(*x[i]))
+ for x_ik in x_i: out+=pack('f',x_ik)
+ print >>file,out
+ file.close()
+
+#---- Data related
+
+def x2twiss(x):
+ '''
+ Calculate twiss from x. Note sig = sqrt(bet*eps).
+
+ Input : x (Nptcl,6)
+ Output: eps,bet,alf,gam
+
+ 2015.07.30
+ '''
+
+ x =[x_i-mean(x_i) for x_i in array(x).transpose()]
+ sig=[[mean(x_i*x_k) for x_k in x] for x_i in x]
+ eps=[ sqrt(det(array(sig)[i:i+2][:,i:i+2])) for i in (0,2,4)]
+ bet=[ sig[i ][i ]/eps[i/2] for i in (0,2,4)]
+ alf=[-sig[i ][i+1]/eps[i/2] for i in (0,2,4)]
+ gam=[ sig[i+1][i+1]/eps[i/2] for i in (0,2,4)]
+
+ return eps,bet,alf,gam
+
+def x2twiss_halo(x,sig_cut=(None,None,None)):
+ '''
+ Calculate twiss of halo particles (r > sig_cut).
+ Note halo particles are different for each plane.
+
+ Input : x (Nptcl,6), sig_cut
+ Output: eps,bet,alf,gam
+
+ 2015.07.30
+ '''
+
+ eps,bet,alf,gam=x2twiss(x)
+ for k in (0,2,4):
+ if sig_cut[k/2]!=None:
+ x=x.transpose();
+ r_nom_sq=(x[k]**2+(alf[k/2]*x[k]+bet[k/2]*x[k+1])**2)/(bet[k/2]*eps[k/2])
+ x=x.transpose();
+ x_halo=[x[i] for i in range(len(x)) if r_nom_sq[i]>sig_cut[k/2]**2]
+ eps_tmp,bet_tmp,alf_tmp,gam_tmp=x2twiss(x_halo)
+ eps[k/2]=eps_tmp[k/2]
+ bet[k/2]=bet_tmp[k/2]
+ alf[k/2]=alf_tmp[k/2]
+ gam[k/2]=gam_tmp[k/2]
+
+ return eps,bet,alf,gam
+
+def loss_elem2den(s,loss,path_file_dt,dlt_dt=5e-6):
+ '''
+ This function calculates loss density [W/m] from s, loss,
+ and the file of the DTL's cell and DT lengths.
+
+ 2015.01.30
+ '''
+ # Define l
+ l=[s[i+1]-s[i] for i in range(len(s)-1)]; l.insert(0,0.0)
+
+ # DTL cell and DT lengths.
+ file=open(path_file_dt)
+ l_cell,l_dt=array([map(float,lin.split()) for lin in file.readlines()]).transpose()[:2]
+ file.close()
+
+ # Replace l's in DTL with DT lengths.
+ Ndt=0
+ for i in range(len(l)):
+ dl=list(abs(l_cell-l[i])); dl_min=min(dl)
+ if dl_min<dlt_dt:
+ l[i]=l_dt[dl.index(dl_min)]; Ndt+=1
+
+ # Check the replacement.
+ if Ndt!=len(l_cell):
+ print 'drift-tube # unmatched, in file: '+str(len(l_cell))+', from matching: '+str(Ndt)
+ print 'review the threshold, exiting...'; exit()
+
+ # Treat inf loss den.
+ for i in range(len(l)):
+ 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()
+ else:
+ k=1
+ while l[i-k]==0.0: k+=1
+ loss[i-k]+=loss[i]; loss[i]=0.0
+ print 'Caution: Inf loss density at elem #',i,'!!!!'
+
+ # Finally, convert to loss density.
+ loss_den=[]
+ for i in range(len(l)):
+ if loss[i]==0.0: loss_den.append(0.0)
+ else : loss_den.append(loss[i]/l[i])
+
+ return loss_den
+
+def partran_end_all(file_name_in_all,file_name_out):
+ '''
+ - Reads multiple partran1.out files and extracts the data at the end.
+ - The same format as partran1.out allows to be used by PARTRAN class.
+ - The elem # is replaced with the file #.
+
+ 2015.11.25
+ '''
+
+ # Extract header from 1 file.
+ with open(file_name_in_all[0],'r') as file:
+ header=''
+ for lin in file.readlines():
+ header+=lin
+ if lin.split()[0]=='####': break
+
+ # Extract output data for each file and write.
+ with open(file_name_out,'w') as file_out:
+ file_out.write(header)
+ for i in range(len(file_name_in_all)):
+ lin=check_output(['tail','-1',file_name_in_all[i]]).split()
+ # Replace elem # with file #.
+ file_out.write('%5s'%str(i))
+ # Just to maintain the same format as partran1.out. Maybe too much...
+ for i in range(1,len(lin)):
+ if '.' in lin[i]: file_out.write(' %13s'%lin[i])
+ else : file_out.write(' %10s'%lin[i])
+ file_out.write('\n')
+
+#-------- Obsolete classes and functions
+
+#---- Old MATCH and LATTICE classes from the time of IPAC'15 (more complex)
+
+## class MATCH:
+## '''
+## Class for matching. Could be revised.
+
+## 2015.04.26
+## '''
+## def __init__(self,typ):
+
+## # Instances
+## self.typ =typ # Need typ_knob and etc ??
+## self.i_diag=[]
+## self.i_knob=[]
+
+## # Instances for 'TRAJ'
+## self.Rx_inv=None
+## self.Ry_inv=None
+
+## class LATTICE:
+## '''
+## Class to handle a TraceWin lattice file.
+
+## - For now, only for the steerer and BPM. Extend as needed.
+## - For now, only for the MEBT. Extend as needed.
+## - For now, reading a file for ken. Implement a method for future ??
+
+## 2015.04.12
+## '''
+## def __init__(self,path_file_lat,path_file_ken):
+
+## # Consts, need to define globally ??
+## mass=938.2723; clight=2.99792458
+
+## # Some dict and list (MEBT and DTL are supported for now...)
+## dic_class={'QUAD' : QUAD ,
+## 'THIN_STEERING' : THIN_STEERING ,
+## 'STEERER' : STEERER ,
+## 'DIAG_POSITION' : DIAG_POSITION ,
+## 'ADJUST' : ADJUST ,
+## 'ADJUST_STEERER': ADJUST_STEERER}
+## lst_elem =['DRIFT','QUAD','GAP','DTL_CEL',
+## 'THIN_STEERING','DIAG_POSITION']
+## lst_comm =['ADJUST','ADJUST_STEERER']
+## lst_diag =['DIAG_POSITION']
+
+## # Load ken from TW's "Energy(MeV).txt" (vs Nelem, starting from 0)
+## with open(path_file_ken,'r') as file:
+## ken=[]
+## for lin in file:
+## try : ken.append(float(lin.split()[2]))
+## except: pass
+## ken.insert(0,ken[0])
+
+## # Go through the lat file
+## with open(path_file_lat) as file:
+## lst=[]; i=0; Nelem=0
+## for lin in file:
+## lin=lin.partition(';')[0] # Remove comment
+## if lin.split()!=[]: # Remove empty line
+## # Convert an elem/comm to a class (maybe too cryptic...)
+## try : lst.append(dic_class[ALLTYPE(lin).typ](lin))
+## except: lst.append( ALLTYPE(lin))
+## # Assign index
+## lst[-1].indx=i; i+=1
+## # Assign Nelem
+## if lst[-1].typ in lst_elem: Nelem+=1
+## lst[-1].Nelem=Nelem
+## # Assign gamma, beta, Brho
+## gamma=1.0+ken[Nelem]/mass ; lst[-1].gamma=gamma
+## beta =sqrt(1.0-1.0/gamma**2) ; lst[-1].beta =beta
+## Brho =(10.0/clight)*gamma*beta*mass*1e-3; lst[-1].Brho =Brho
+## # End at 'END'
+## if lst[-1].typ=='END': break
+
+## # Go through lat again and define matchings, extend/modify as needed
+## match={}
+## for i in range(len(lst)):
+## if lst[i].typ in lst_comm:
+## lst,match=lst[i].activate(lst,match)
+
+## # Get index of diagnostics
+## for i in range(len(lst)):
+## if lst[i].typ in lst_diag:
+## try : match[lst[i].fam].i_diag.append(i)
+## except: pass
+
+## # Instances
+## self.lst =lst
+## self.match=match
+
+## def get_tw(self,path_file,flag_no_err=None):
+
+## with open(path_file,'w') as file:
+## for lst_i in self.lst:
+## if 'ERROR' in lst_i.typ:
+## if flag_no_err!=None: pass
+## else : file.write(lst_i.get_tw())
+## else:
+## file.write(lst_i.get_tw())
+
+#---- Old DENSITY class for the envelope calc only
+
+## class DENSITY_ENV:
+## '''
+## This class is to handle Density_Env.dat.
+
+## - For now, this is intended for an individual file and NOT "tot" file.
+## (Something not clear between an individual file and "tot" file...)
+## - For now, there are instances of only Nelem, s, x_ave, y_ave.
+## (Add others when needed.)
+## - For now, tested only for the MEBT.
+## - Note there are much more data points than the total number of elems
+## due to the slicing of the SC calculation. Use i_elem to extract data
+## at the ends of elems, e.g., array(x_ave)[array(i_elem)].
+
+## 2015.03.29
+## '''
+## def __init__(self,path_file):
+
+## # Go through the file
+## Nelem=[]; s=[]; x=[]
+## with open(path_file) as file:
+## while True:
+## try:
+## ver=fromfile(file,dtype=uint16,count=3)[0] # version, year, vlong
+## Nelem.append(fromfile(file,dtype=uint32,count=2)[1]) # Nrun, Nelem
+## s.append(fromfile(file,dtype=float32,count=4)[1]) # Ibeam, s, aptx, apty
+## fromfile(file,dtype=uint32,count=1) # Nslice ??
+## x.append(list(fromfile(file,dtype=float32,count=7*4)[:2])) # ave & etc
+## if ver>=5: fromfile(file,dtype=float32,count=7*2) # ver >= 5 part
+## if ver>=6: fromfile(file,dtype=float32,count=7*2) # ver >= 6 part
+## if ver>=7: fromfile(file,dtype=float32,count=3*2) # ver >= 7 part
+## if ver>=8: fromfile(file,dtype=float32,count=1*3) # ver >= 8 part
+## fromfile(file,dtype=uint64,count=1) # Nptcl
+## except:
+## break
+
+## # Index of n-th elem end, starting from 0
+## i_elem=[0]
+## for n in range(1,Nelem[-1]):
+## k=1
+## while True:
+## try : i_elem.append(Nelem.index(n+k)-1); break
+## except: k+=1
+## i_elem.append(len(Nelem)-1)
+
+## # Instances
+## self.Nelem=Nelem
+## self.s =s
+## self.x =list(1e3*array(x).transpose()[0])
+## self.y =list(1e3*array(x).transpose()[1])
+
+## # Additional instances
+## self.i_elem=i_elem
+
+#---- Lattice related functions
+
+## def clean_lat(path_file):
+
+## '''
+## Remove comments, names, and etc from a TraceWin file.
+## To convert the output to str, use: [' '.join(l)+'\n' for l in lat].
+## The current version not capitalizing.
+
+## Output: Lattice in a list
+
+## 2015.03.18
+## '''
+
+## with open(path_file) as file:
+## lat=[]
+## for lin in file:
+## lin=lin.partition(';')[0] # Remove comment
+## #lin=lin.upper().partition(';')[0] # Remove comment and capitalize
+## if ':' in lin: lin=lin.partition(':')[-1].split() # Remove name
+## else : lin=lin.split()
+## if lin!=[]: # Avoid empty line
+## lat.append(lin)
+## if lin[0].upper()=='END': break
+
+## return lat
+
+## def get_steerer(path_file):
+
+## '''
+## Extract steerer values from "Adjusted_Values.txt" of TraceWin.
+## Note the key is the "lat elem #" where commands are not counted.
+
+## Output is dic of {lat elem #:[val]} or {lat elem #:[val_x,val_y]}
+
+## 2014.10.17
+## '''
+
+## file=open(path_file)
+## i_elem=0; dic_steerer_val={}
+## for lin in file.readlines():
+## lin=lin.split()
+## for i in range(len(lin)):
+## if lin[i]==':':
+## if int(lin[i-1])!=i_elem:
+## i_elem=int(lin[i-1])
+## dic_steerer_val[i_elem]=[lin[i+1]]
+## else:
+## dic_steerer_val[i_elem].append(lin[i+1])
+## file.close()
+
+## return dic_steerer_val
+
+## def apply_steerer(lat,dic_steerer_val,lst_typ_elem):
+
+## '''
+## Apply steerer values to a TraceWin lattice.
+
+## Input: lat from "clean_lat, steerer dic from "get_steerer", list of lat elem types
+
+## 2014.10.17
+## '''
+
+## #--
+## i_elem=0; dic_steerer={}
+## for i in range(len(lat)):
+## if lat[i][0] in lst_typ_elem:
+## i_elem+=1
+## if lat[i][0]=='STEERER':
+## if i_elem+1 in dic_steerer_val.keys():
+## if len(dic_steerer_val[i_elem+1])==2:
+## dic_steerer[i]=[[1,dic_steerer_val[i_elem+1][0]],[2,dic_steerer_val[i_elem+1][1]]]
+## else:
+## for k in range(i)[::-1]:
+## if lat[k][0]=='ADJUST_STEERER_BX':
+## dic_steerer[i]=[[1,dic_steerer_val[i_elem+1][0]]]; break
+## if lat[k][0]=='ADJUST_STEERER_BY':
+## dic_steerer[i]=[[2,dic_steerer_val[i_elem+1][0]]]; break
+## if lat[i][0]=='THIN_STEERING':
+## if i_elem in dic_steerer_val.keys():
+## if len(dic_steerer_val[i_elem])==2:
+## dic_steerer[i]=[[1,dic_steerer_val[i_elem][0]],[2,dic_steerer_val[i_elem][1]]]
+## else:
+## for k in range(i)[::-1]:
+## if lat[k][0]=='ADJUST' and lat[k][2]=='1':
+## dic_steerer[i]=[[1,dic_steerer_val[i_elem][0]]]; break
+## if lat[k][0]=='ADJUST' and lat[k][2]=='2':
+## dic_steerer[i]=[[2,dic_steerer_val[i_elem][0]]]; break
+
+## for i in sorted(dic_steerer.keys()):
+## for k in range(len(dic_steerer[i])):
+## lat[i][dic_steerer[i][k][0]]=dic_steerer[i][k][1]
+
+## return lat
+
+## def apply_multipole(lat,typ,r,b3=0.0,b4=0.0,b5=0.0,b6=0.0):
+
+## '''
+## Apply multipole components to quads in a TraceWin lattice file.
+## Note all the quads in all the sections are affected in the same way.
+
+## Input : lat from "clean_lat, reference r, b3-b6, and dist type
+## ("fix", "uniform", "gauss")
+
+## 2015.01.08
+## '''
+
+## bn=(b3,b4,b5,b6)
+
+## for i in range(len(lat)):
+## if lat[i][0]=='QUAD':
+
+## # Adjust format
+## if len(lat[i])<5:
+## for k in range(5): lat[i].append('0')
+## else:
+## lat[i]=lat[i][:5]
+## for k in range(4): lat[i].append('0')
+
+## # Fixed value
+## if typ=='fix':
+## for k in range(4):
+## lat[i][k+5]=str( bn[k]/r**(k+1)*float(lat[i][2]))
+## # Uniform dist
+## if typ=='uniform':
+## for k in range(4):
+## lat[i][k+5]=str((2.0*rand()-1.0)*bn[k]/r**(k+1)*float(lat[i][2]))
+## # Gaussian dist
+## if typ=='gauss':
+## for k in range(4):
+## lat[i][k+5]=str( randn()*bn[k]/r**(k+1)*float(lat[i][2]))
+
+## return lat
+
+#---- Dist related
+
+## def partran2twiss(path_file):
+
+## '''
+## MERGE THIS TO PARTRAN CLASS !!!!
+
+## Extract output Twiss and halo from partran1.out
+
+## Longitudinal phase space is z-z'.
+
+## 2015.01.15
+## '''
+
+## file=open(path_file)
+
+## for lin in file.readlines():
+## lin=lin.split()
+## if lin[0]=='####':
+## idx_gamma=lin.index('gama-1')
+## idx_eps =[lin.index(p) for p in ("ex" ,"ey" ,"ezdp" )]
+## idx_bet =[lin.index(p) for p in ("SizeX","SizeY","SizeZ")]
+## idx_alf =[lin.index(p) for p in ("sxx'" ,"syy'" ,"szdp" )]
+## idx_hal =[lin.index(p) for p in ("hx" ,"hy" ,"hp" )]
+## file.close()
+
+## gamma=float(lin[idx_gamma])+1.0
+## eps =array([float(lin[idx]) for idx in idx_eps])
+## bet =array([float(lin[idx])**2 for idx in idx_bet]) # Actually Sig_{i,i}
+## alf =array([float(lin[idx]) for idx in idx_alf]) # Actually Sig_{i,i+1}
+## hal =array([float(lin[idx]) for idx in idx_hal])
+
+## bet= bet/eps*sqrt(gamma**2-1.0); bet[2]*=gamma**2 # z-dp => z-z'
+## alf=-alf/eps*sqrt(gamma**2-1.0)
+## gam= (1.0+alf**2)/bet
+
+## return eps,bet,alf,gam,hal
+
+## class PROJECT_SING_CORE:
+## '''
+## - Use this for simultaneous mult 1-core-simulations (w/ tracelx64).
+## - Initial parameters are not meant to be changed, unlike "PROJECT".
+## (Each run should have its own project.)
+## - Make sure the project/lattice file is in the calc dir.
+## - Changing the calc dir option not working for tracelx64.
+
+## 2015.10.15
+## '''
+
+## def __init__(self,path_cal,
+## file_name_proj='proj.ini',file_name_lat='lat.dat',seed=None,flag_hide=None):
+
+## # Make sure "path_cal" ends with "/".
+## if path_cal[-1]!='/': path_cal+='/'
+
+## # Check the calc dir exists.
+## if isdir(path_cal)==False: print 'Calc dir does not exist !!!! Exiting ...'; exit()
+
+## # Check "tracelx64" is in the calc dir.
+## if isfile(path_cal+'tracelx64')==False:
+## try : system('cp /opt/cea/tracelx64 '+path_cal)
+## except: print 'tracelx64 not in calc dir !!!! Exiting ...'; exit()
+
+## # Check "tracewin.key" is in the calc dir.
+## if isfile(path_cal+'tracewin.key')==False:
+## try : system('cp /opt/cea/tracewin.key '+path_cal)
+## except: print 'tracewin.key not in calc dir !!!! Exiting ...'; exit()
+
+## # Check the project/lattice file is in the calc dir.
+## if '/' not in file_name_proj: file_name_proj=path_cal+file_name_proj
+## if '/' not in file_name_lat : file_name_lat =path_cal+file_name_lat
+## if isfile(file_name_proj)==False: print 'project file not in calc dir !!!! Exiting ...'; exit()
+## if isfile(file_name_lat) ==False: print 'lattice file not in calc dir !!!! Exiting ...'; exit()
+
+## # Instances (Add more options as needed.)
+## self.path_cal =path_cal
+## self.file_name_proj=file_name_proj
+## self.file_name_lat =file_name_lat
+## self.seed =seed
+## self.flag_hide =flag_hide
+
+## def exe(self):
+
+## opt_exe =self.path_cal+'tracelx64 '+self.file_name_proj
+## if self.file_name_lat!=None: opt_exe+=' dat_file=' +self.file_name_lat
+## if self.seed !=None: opt_exe+=' random_seed='+self.seed
+## if self.flag_hide !=None: opt_exe+=' hide'
+
+## system(opt_exe)
diff --git a/ess/lib_tw_elem.py b/ess/lib_tw_elem.py
new file mode 100644
index 0000000000000000000000000000000000000000..fe96ea43346c8692536fc122842e238a117a6ed4
--- /dev/null
+++ b/ess/lib_tw_elem.py
@@ -0,0 +1,1049 @@
+
+'''
+ - Classes for TraceWin elements/commands for external manipulations.
+ - First double check all the used elements/commands.
+
+ - Element/command types
+
+ * Active elements
+
+ FIELD_MAP
+
+ QUAD
+ THIN_STEERING
+ GAP
+ DTL_CEL
+
+ * Passive elements
+
+ DRIFT
+ BEND
+ EDGE
+ APERTURE
+ DIAG_POSITION
+ DIAG_...
+
+ * Active commands
+
+ STEERER
+ CHOPPER
+
+ * Passive commands
+
+ FREQ
+ MARKER
+
+ * Other commands for just common "COMM"
+
+ END
+ LATTICE
+ LATTICE_END
+
+ SET_ADV
+ SET_SIZE
+ SET_TWISS
+ SET_ACHROMAT
+ MATCH_FAM_GRAD
+ MATCH_FAM_FIELD
+ MATCH_FAM_PHASE
+ MATCH_FAM_LFOC
+ MIN_PHASE_VARIATION
+ START_ACHROMAT
+
+ ERROR_QUAD_NCPL_STAT
+ ERROR_QUAD_CPL_STAT
+ ERROR_CAV_NCPL_STAT
+ ERROR_CAV_NCPL_DYN
+ ERROR_CAV_CPL_STAT
+ ERROR_CAV_CPL_DYN
+ ADJUST
+ ADJUST_STEERER
+ ADJUST_STEERER_BX
+ ADJUST_STEERER_BY
+ SET_BEAM_PHASE_ERROR
+
+ PLOT_DST
+
+ 2015.10.14
+'''
+
+#---- Libs
+
+from numpy import *
+from sys import exit
+
+#---- Constants
+
+#mass=938.272046 # Wiki
+mass=938.272029 # TraceWin
+c =2.99792458
+
+#---- Misc functions/classes
+
+def Brho(gamma):
+
+ return (10.0/c)*sqrt(gamma**2-1.0)*mass*1e-3
+
+class FIELD_MAP_DATA:
+ '''
+ - Only 1D (Ez) supported for now.
+ '''
+ def __init__(self,file_name):
+
+ # Read file
+ with open(file_name) as file: lin=file.readlines()
+
+ # Instances
+ self.dL =float(lin[0].split()[1])/int(lin[0].split()[0])
+ self.field=map(float,lin[2:])
+
+#---- Parent classes
+
+class ELEM:
+ '''
+ '''
+ def __init__(self,name,typ,para):
+
+ # Basic instances
+ self.name=name
+ self.typ =typ
+ self.para=para
+
+ # Option instances
+ self.L = 0.0
+ self.apt = None
+ self.idx =-1
+ self.idx_elem=-1
+ self.s = 0.0
+ self.gamma = 1.0
+ self.freq = 352.21
+
+ def update_idx(self):
+
+ # Update idx, idx_elem, s
+ self.idx +=1
+ self.idx_elem+=1
+ self.s +=self.L
+
+ def get_tw(self):
+
+ if self.name!='': lin=self.name+': '+self.typ+' '+' '.join(self.para)
+ if self.name=='': lin= self.typ+' '+' '.join(self.para)
+
+ return lin
+
+ ## def get_madx(self):
+
+ ## i='0'*(3-int(log10(self.idx_elem+1)))+str(self.idx_elem+1)
+
+ ## if self.name!='': lin= self.name +': marker;'
+ ## if self.name=='': lin='elem'+i+'_marker'+': marker;'
+
+ ## return lin
+
+class COMM:
+ '''
+ - Same inputs as ELEM class but some may not be needed ...
+ '''
+ def __init__(self,name,typ,para):
+
+ # Basic instances
+ self.name=name
+ self.typ =typ
+ self.para=para
+
+ # Option instances
+ self.idx =-1
+ self.idx_elem=-1
+ self.s = 0.0
+ self.gamma = 1.0
+ self.freq = 352.21
+
+ def update_idx(self):
+
+ # Update idx
+ self.idx+=1
+
+ def get_tw(self):
+
+ if self.name!='': lin=self.name+': '+self.typ+' '+' '.join(self.para)
+ if self.name=='': lin= self.typ+' '+' '.join(self.para)
+
+ return lin
+
+#---- Field map (need an additional input, dic_fmap)
+
+class FIELD_MAP(ELEM):
+ '''
+ - Calculation of gamma only 1D (Ez) supported.
+ '''
+ def __init__(self,name,typ,para,dic_fmap):
+
+ ELEM.__init__(self,name,typ,para)
+
+ # TW instances
+ self.name_fmap= para[8]
+ self.typ_fmap = para[0]
+ self.L =float(para[1])*1e-3 # mm => m
+ self.phs_rf =float(para[2])*pi/180.0 # deg => rad
+ self.apt =float(para[3])
+ self.Bnom =float(para[4])
+ self.Enom =float(para[5])
+ self.sc_comp =float(para[6])
+ self.flag_apt = para[7]
+
+ # Option instances
+ self.data=dic_fmap[self.name_fmap]
+
+ def update_gamma(self):
+
+ # Temp error message for field map types
+ if self.typ_fmap!='100': print 'Gamma calc only supported for 1D, exiting ...', exit()
+
+ # Update gamma (simple ver, closer to TW)
+ dL=self.data.dL; V=self.Enom*array(self.data.field)*dL; phs=self.phs_rf; C=0.0; S=0.0
+ for i in range(len(V)):
+ C+=V[i]*cos(phs); S+=V[i]*sin(phs); self.gamma+=V[i]/mass*cos(phs)
+ beta=sqrt(1.0-1.0/self.gamma**2); phs+=2.0*pi*self.freq*dL/(beta*c*1e2)
+
+ # Update gamma (I think this is more accurate...)
+ ## dL=self.data.dL; V=self.Enom*array(self.data.field)*dL; phs=self.phs_rf; C=0; S=0
+ ## for i in range(len(V)):
+ ## if i==0 :
+ ## C+=V[i]*cos(phs); S+=V[i]*sin(phs); self.gamma+=0.5*V[i]/mass*cos(phs)
+ ## beta=sqrt(1.0-1.0/self.gamma**2); phs+=2.0*pi*self.freq*dL/(beta*c*1e2)
+ ## if 0<i<len(V)-1:
+ ## C+=V[i]*cos(phs); S+=V[i]*sin(phs); self.gamma+= V[i]/mass*cos(phs)
+ ## beta=sqrt(1.0-1.0/self.gamma**2); phs+=2.0*pi*self.freq*dL/(beta*c*1e2)
+ ## if i==len(V)-1 :
+ ## C+=V[i]*cos(phs); S+=V[i]*sin(phs); self.gamma+=0.5*V[i]/mass*cos(phs)
+
+ # Option instances for other codes (MADX?)
+ self.phs_syn=arctan(S/C)
+ self.E0TL =C/cos(self.phs_syn)
+
+ def get_tw(self):
+
+ if self.name!='': lin=self.name+': '+self.typ+' '
+ if self.name=='': lin= self.typ+' '
+ lin+= self.typ_fmap +' '
+ lin+=str(self.L*1e3 )+' '
+ lin+=str(self.phs_rf*180.0/pi)+' '
+ lin+=str(self.apt )+' '
+ lin+=str(self.Bnom )+' '
+ lin+=str(self.Enom )+' '
+ lin+=str(self.sc_comp )+' '
+ lin+= self.flag_apt +' '
+ lin+= self.name_fmap
+
+ return lin
+
+ def get_madx(self):
+
+ i='0'*(3-int(log10(self.idx_elem+1)))+str(self.idx_elem+1)
+
+ if self.name!='': lin= self.name +': RFCAVITY, '
+ if self.name=='': lin='ELEM'+i+'_FIELDMAP'+': RFCAVITY, '
+ lin+='L=' +str(self.L )+', '
+ lin+='FREQ='+str(self.freq )+', '
+ lin+='VOLT='+str(self.E0TL )+', '
+ lin+='LAG=' +str((0.5*pi-self.phs_syn)/(2.0*pi))+'; ' # phs_MADX = pi/2-phs_TW
+
+ return lin
+
+ def get_fluka(self):
+
+ i='0'*(3-int(log10(self.idx_elem+1)))+str(self.idx_elem+1)
+
+ if self.name!='': lin= self.name +' RFCAVITY '
+ if self.name=='': lin='ELEM'+i+'_FIELDMAP'+' RFCAVITY '
+ lin+=str(self.L)+' '+str(self.s)+' '
+ lin+='0.0 0.0 0.0 '
+ lin+='CIRCLE '+str(self.apt*1e-3)+' 0.0'
+
+ return lin
+
+ def get_mars(self):
+
+ lin ='"RFCAVITY" "" "" '+str(self.s)+' '+str(self.L)+' '
+ lin+='0 0 0 0 0'
+
+ return lin
+
+#---- Active elements
+
+class QUAD(ELEM):
+ '''
+ - Tilt not supported.
+ '''
+ def __init__(self,name,typ,para):
+
+ ELEM.__init__(self,name,typ,para)
+
+ # TW instances
+ self.L =float(para[0])*1e-3 # mm => m
+ self.G =float(para[1])
+ self.apt =float(para[2])
+ self.tilt=0.0
+
+ # TW option instances
+ try : self.b3 =float(para[4])
+ except: self.b3 =0.0
+ try : self.b4 =float(para[5])
+ except: self.b4 =0.0
+ try : self.b5 =float(para[6])
+ except: self.b5 =0.0
+ try : self.b6 =float(para[7])
+ except: self.b6 =0.0
+ try : self.r_good=float(para[8])
+ except: self.r_good=0.0
+
+ def get_tw(self):
+
+ if self.name!='': lin=self.name+': '+self.typ+' '
+ if self.name=='': lin= self.typ+' '
+ lin+=str(self.L*1e3 )+' '
+ lin+=str(self.G )+' '
+ lin+=str(self.apt )+' '
+ lin+=str(self.tilt )+' '
+ lin+=str(self.b3 )+' '
+ lin+=str(self.b4 )+' '
+ lin+=str(self.b5 )+' '
+ lin+=str(self.b6 )+' '
+ lin+=str(self.r_good)
+
+ return lin
+
+ def get_madx(self):
+
+ i='0'*(3-int(log10(self.idx_elem+1)))+str(self.idx_elem+1)
+
+ if self.name!='': lin= self.name +': QUADRUPOLE, '
+ if self.name=='': lin='ELEM'+i+'_QUAD'+': QUADRUPOLE, '
+ lin+='L=' +str(self.L )+', '
+ lin+='K1='+str(self.G/Brho(self.gamma))+'; '
+
+ return lin
+
+ def get_fluka(self):
+
+ i='0'*(3-int(log10(self.idx_elem+1)))+str(self.idx_elem+1)
+
+ if self.name!='': lin= self.name +' QUADRUPOLE '
+ if self.name=='': lin='ELEM'+i+'_QUAD'+' QUADRUPOLE '
+ lin+=str(self.L)+' '+str(self.s)+' '
+ lin+='0.0 0.0 '+str(self.G*self.L/Brho(self.gamma))+' '
+ lin+='CIRCLE '+str(self.apt*1e-3)+' 0.0'
+
+ return lin
+
+ def get_mars(self):
+
+ lin ='"QUADRUPOLE" "" "" '+str(self.s)+' '+str(self.L)+' '
+ lin+='0 '+str(self.G/Brho(self.gamma))+' 0 0 0 '
+
+ return lin
+
+class THIN_STEERING(ELEM):
+ '''
+ - Only magnetic steerer supported.
+ - Info of plane missing.
+ '''
+ def __init__(self,name,typ,para):
+
+ ELEM.__init__(self,name,typ,para)
+
+ # TW instances
+ self.BLx =float(self.para[0])
+ self.BLy =float(self.para[1])
+ self.apt =float(self.para[2])
+ self.typ_EM='0'
+
+ def get_tw(self):
+
+ if self.name!='': lin=self.name+': '+self.typ+' '
+ if self.name=='': lin= self.typ+' '
+ lin+=str(self.BLx )+' '
+ lin+=str(self.BLy )+' '
+ lin+=str(self.apt )+' '
+ lin+= self.typ_EM
+
+ return lin
+
+ def get_madx(self):
+
+ i='0'*(3-int(log10(self.idx_elem+1)))+str(self.idx_elem+1)
+
+ if self.name!='': lin= self.name +': KICKER, L=0, '
+ if self.name=='': lin='ELEM'+i+'_STEERER'+': KICKER, L=0, '
+ lin+='HKICK='+str(-self.BLy/Brho(self.gamma))+', '
+ lin+='VKICK='+str( self.BLx/Brho(self.gamma))+'; '
+
+ return lin
+
+ def get_fluka(self):
+
+ i='0'*(3-int(log10(self.idx_elem+1)))+str(self.idx_elem+1)
+
+ if self.name!='': lin= self.name +' KICKER '
+ if self.name=='': lin='ELEM'+i+'_STEERER'+' KICKER '
+ lin+=str(self.L)+' '+str(self.s)+' '
+ lin+='0.0 0.0 0.0 '
+ lin+='CIRCLE '+str(self.apt*1e-3)+' 0.0'
+
+ return lin
+
+ def get_mars(self):
+
+ lin ='"KICKER" "" "" '+str(self.s)+' '+str(self.L)+' '
+ lin+='0 0 0 0 0'
+
+ return lin
+
+class GAP(ELEM):
+ '''
+ - Note there are E0T0L (for TW) and E0TL (for other codes).
+ '''
+ def __init__(self,name,typ,para):
+
+ ELEM.__init__(self,name,typ,para)
+
+ # TW instances
+ self.E0T0L =float(para[0])*1e-6 # V => MV
+ self.phs_rf =float(para[1])*pi/180.0 # deg => rad
+ self.apt =float(para[2])
+ self.phs_rf_typ='0' # Relative phase only for now
+
+ # TW option instances
+ try : self.beta_s=float(para[4])
+ except: self.beta_s=0.0
+ try : self.T0 =float(para[5])
+ except: self.T0 =0.0
+ try : self.T1 =float(para[6])
+ except: self.T1 =0.0
+ try : self.T2 =float(para[7])
+ except: self.T2 =0.0
+
+ # Option instances
+ self.E0TL=self.E0T0L
+
+ def update_gamma(self):
+
+ # Update gamma (and E0TL)
+ if self.beta_s!=0.0 and self.T0!=0.0:
+ gamma_c =self.gamma+0.5*self.E0TL/mass*cos(self.phs_rf)
+ beta_c =sqrt(1.0-1.0/gamma_c**2)
+ k =self.beta_s/beta_c
+ self.E0TL*=(self.T0+self.T1*(k-1.0)+0.5*self.T2*(k-1.0)**2)/self.T0
+ self.gamma+=self.E0TL/mass*cos(self.phs_rf)
+
+ def get_tw(self):
+
+ if self.name!='': lin=self.name+': '+self.typ+' '
+ if self.name=='': lin= self.typ+' '
+ lin+=str(self.E0T0L*1e6 )+' '
+ lin+=str(self.phs_rf*180.0/pi)+' '
+ lin+=str(self.apt )+' '
+ lin+= self.phs_rf_typ +' '
+ lin+=str(self.beta_s )+' '
+ lin+=str(self.T0 )+' '
+ lin+=str(self.T1 )+' '
+ lin+=str(self.T2 )
+
+ return lin
+
+ def get_madx(self):
+
+ i='0'*(3-int(log10(self.idx_elem+1)))+str(self.idx_elem+1)
+
+ if self.name!='': lin= self.name +': RFCAVITY, L=0, '
+ if self.name=='': lin='ELEM'+i+'_GAP'+': RFCAVITY, L=0, '
+ lin+='FREQ='+str(self.freq )+', '
+ lin+='VOLT='+str(self.E0TL )+', '
+ lin+='LAG=' +str((0.5*pi-self.phs_rf)/(2.0*pi))+'; ' # phs_MADX = pi/2-phs_TW
+
+ return lin
+
+ def get_fluka(self):
+
+ i='0'*(3-int(log10(self.idx_elem+1)))+str(self.idx_elem+1)
+
+ if self.name!='': lin= self.name +' RFCAVITY '
+ if self.name=='': lin='ELEM'+i+'_GAP'+' RFCAVITY '
+ lin+=str(self.L)+' '+str(self.s)+' '
+ lin+='0.0 0.0 0.0 '
+ lin+='CIRCLE '+str(self.apt*1e-3)+' 0.0'
+
+ return lin
+
+ def get_mars(self):
+
+ lin ='"RFCAVITY" "" "" '+str(self.s)+' '+str(self.L)+' '
+ lin+='0 0 0 0 0'
+
+ return lin
+
+class DTL_CEL(ELEM):
+ '''
+ - Note there are E0T0L (for TW) and E0TL (for other codes).
+ '''
+ def __init__(self,name,typ,para):
+
+ ELEM.__init__(self,name,typ,para)
+
+ # TW instances
+ self.L =float(para[0])*1e-3 # mm => m
+ self.L_Q1 =float(para[1])*1e-3 # mm => m
+ self.L_Q2 =float(para[2])*1e-3 # mm => m
+ self.ds_gap =float(para[3])*1e-3 # mm => m
+ self.G_Q1 =float(para[4])
+ self.G_Q2 =float(para[5])
+ self.E0T0L =float(para[6])*1e-6 # V => MV
+ self.phs_rf =float(para[7])*pi/180.0 # deg => rad
+ self.apt =float(para[8])
+ self.phs_rf_typ='0' # Relative phase only for now
+
+ # TW option instances
+ try : self.beta_s=float(para[10])
+ except: self.beta_s=0.0
+ try : self.T0 =float(para[11])
+ except: self.T0 =0.0
+ try : self.T1 =float(para[12])
+ except: self.T1 =0.0
+ try : self.T2 =float(para[13])
+ except: self.T2 =0.0
+
+ # Option instances
+ self.E0TL=self.E0T0L
+
+ def update_gamma(self):
+
+ # Save the gamma_ini (for MADX)
+ self.gamma_ini=self.gamma
+
+ # Update gamma (and E0TL)
+ if self.beta_s!=0.0 and self.T0!=0.0:
+ gamma_c =self.gamma+0.5*self.E0TL/mass*cos(self.phs_rf)
+ beta_c =sqrt(1.0-1.0/gamma_c**2)
+ k =self.beta_s/beta_c
+ self.E0TL*=(self.T0+self.T1*(k-1.0)+0.5*self.T2*(k-1.0)**2)/self.T0
+ self.gamma+=self.E0TL/mass*cos(self.phs_rf)
+
+ def get_tw(self):
+
+ if self.name!='': lin=self.name+': '+self.typ+' '
+ if self.name=='': lin= self.typ+' '
+ lin+=str(self.L*1e3 )+' '
+ lin+=str(self.L_Q1*1e3 )+' '
+ lin+=str(self.L_Q2*1e3 )+' '
+ lin+=str(self.ds_gap*1e3 )+' '
+ lin+=str(self.G_Q1 )+' '
+ lin+=str(self.G_Q2 )+' '
+ lin+=str(self.E0T0L*1e6 )+' '
+ lin+=str(self.phs_rf*180.0/pi)+' '
+ lin+=str(self.apt )+' '
+ lin+= self.phs_rf_typ +' '
+ lin+=str(self.beta_s )+' '
+ lin+=str(self.T0 )+' '
+ lin+=str(self.T1 )+' '
+ lin+=str(self.T2 )
+
+ return lin
+
+ def get_madx(self):
+
+ i='0'*(3-int(log10(self.idx_elem+1)))+str(self.idx_elem+1)
+
+ if self.name!='': lin= self.name +'_QUAD1: QUADRUPOLE, '
+ if self.name=='': lin='ELEM'+i+'_DTLCELL'+'_QUAD1: QUADRUPOLE, '
+ lin+='L=' +str(self.L_Q1 )+', '
+ lin+='K1='+str(self.G_Q1/Brho(self.gamma_ini))+'; \n'
+
+ if self.name!='': lin+= self.name +'_DRIFT1: DRIFT, '
+ if self.name=='': lin+='ELEM'+i+'_DTLCELL'+'_DRIFT1: DRIFT, '
+ lin+='L='+str(0.5*self.L-self.L_Q1-self.ds_gap)+';\n'
+
+ if self.name!='': lin+= self.name +'_GAP: RFCAVITY, L=0, '
+ if self.name=='': lin+='ELEM'+i+'_DTLCELL'+'_GAP: RFCAVITY, L=0, '
+ lin+='FREQ='+str(self.freq )+', '
+ lin+='VOLT='+str(self.E0TL )+', '
+ lin+='LAG=' +str((0.5*pi-self.phs_rf)/(2.0*pi))+'; \n' # phs_MADX = pi/2-phs_TW
+
+ if self.name!='': lin+= self.name +'_DRIFT2: DRIFT, '
+ if self.name=='': lin+='ELEM'+i+'_DTLCELL'+'_DRIFT2: DRIFT, '
+ lin+='L='+str(0.5*self.L-self.L_Q2+self.ds_gap)+';\n'
+
+ if self.name!='': lin+= self.name +'_QUAD2: QUADRUPOLE, '
+ if self.name=='': lin+='ELEM'+i+'_DTLCELL'+'_QUAD2: QUADRUPOLE, '
+ lin+='L=' +str(self.L_Q2 )+', '
+ lin+='K1='+str(self.G_Q2/Brho(self.gamma))+'; '
+
+ return lin
+
+ def get_fluka(self):
+
+ i='0'*(3-int(log10(self.idx_elem+1)))+str(self.idx_elem+1)
+
+ ## if self.name!='': lin= self.name +'_QUAD1 QUADRUPOLE '
+ ## if self.name=='': lin='ELEM'+i+'_DTLCELL'+'_QUAD1 QUADRUPOLE '
+ ## lin+=str(self.L_Q1)+' '+str(self.s-self.L+self.L_Q1)+' '
+ ## lin+='0.0 0.0 '+str(self.G_Q1*self.L_Q1/Brho(self.gamma_ini))+' '
+ ## lin+='CIRCLE '+str(self.apt*1e-3)+' 0.0\n'
+
+ ## if self.name!='': lin+= self.name +'_DRIFT1 DRIFT '
+ ## if self.name=='': lin+='ELEM'+i+'_DTLCELL'+'_DRIFT1 DRIFT '
+ ## lin+=str(0.5*self.L-self.L_Q1-self.ds_gap)+' '+str(self.s-0.5*self.L-self.ds_gap)+' '
+ ## lin+='0.0 0.0 0.0 '
+ ## lin+='CIRCLE '+str(self.apt*1e-3)+' 0.0\n'
+
+ ## if self.name!='': lin+= self.name +'_GAP RFCAVITY '
+ ## if self.name=='': lin+='ELEM'+i+'_DTLCELL'+'_GAP RFCAVITY '
+ ## lin+='0.0 '+str(self.s-0.5*self.L-self.ds_gap)+' '
+ ## lin+='0.0 0.0 0.0 '
+ ## lin+='CIRCLE '+str(self.apt*1e-3)+' 0.0\n'
+
+ ## if self.name!='': lin+= self.name +'_DRIFT2 DRIFT '
+ ## if self.name=='': lin+='ELEM'+i+'_DTLCELL'+'_DRIFT2 DRIFT '
+ ## lin+=str(0.5*self.L-self.L_Q2+self.ds_gap)+' '+str(self.s-self.L_Q2)+' '
+ ## lin+='0.0 0.0 0.0 '
+ ## lin+='CIRCLE '+str(self.apt*1e-3)+' 0.0\n'
+
+ ## if self.name!='': lin+= self.name +'_QUAD2 QUADRUPOLE '
+ ## if self.name=='': lin+='ELEM'+i+'_DTLCELL'+'_QUAD2 QUADRUPOLE '
+ ## lin+=str(self.L_Q2)+' '+str(self.s)+' '
+ ## lin+='0.0 0.0 '+str(self.G_Q2*self.L_Q2/Brho(self.gamma))+' '
+ ## lin+='CIRCLE '+str(self.apt*1e-3)+' 0.0'
+
+ if self.name!='': lin= self.name +' RFCAVITY '
+ if self.name=='': lin='ELEM'+i+'_DTLCELL'+' RFCAVITY '
+ lin+=str(self.L)+' '+str(self.s)+' '
+ lin+='0.0 0.0 0.0 '
+ lin+='CIRCLE '+str(self.apt*1e-3)+' 0.0'
+
+ return lin
+
+ def get_mars(self):
+
+ lin ='"QUADRUPOLE" "" "" '
+ lin+=str(self.s-self.L+self.L_Q1)+' '
+ lin+=str(self.L_Q1) +' '
+ lin+='0 '+str(self.G_Q1/Brho(self.gamma_ini))+' 0 0 0 \n'
+
+ lin+='"DRIFT" "" "" '
+ lin+=str(self.s-0.5*self.L-self.ds_gap) +' '
+ lin+=str(0.5*self.L-self.L_Q1-self.ds_gap)+' '
+ lin+='0 0 0 0 0 \n'
+
+ lin+='"RFCAVITY" "" "" '
+ lin+=str(self.s-0.5*self.L-self.ds_gap)+' '
+ lin+='0 '
+ lin+='0 0 0 0 0 \n'
+
+ lin+='"DRIFT" "" "" '
+ lin+=str(self.s-self.L_Q2) +' '
+ lin+=str(0.5*self.L-self.L_Q2+self.ds_gap)+' '
+ lin+='0 0 0 0 0 \n'
+
+ lin+='"QUADRUPOLE" "" "" '
+ lin+=str(self.s) +' '
+ lin+=str(self.L_Q2)+' '
+ lin+='0 '+str(self.G_Q2/Brho(self.gamma))+' 0 0 0 '
+
+ return lin
+
+#---- Passive elements
+
+class DRIFT(ELEM):
+ '''
+ '''
+ def __init__(self,name,typ,para):
+
+ ELEM.__init__(self,name,typ,para)
+
+ # TW instances
+ self.L =float(para[0])*1e-3 # mm => m
+ self.apt=float(para[1])
+
+ # TW option instances
+ try : self.apty=float(para[2])
+ except: self.apty=0.0
+
+ def get_madx(self):
+
+ i='0'*(3-int(log10(self.idx_elem+1)))+str(self.idx_elem+1)
+
+ if self.name!='': lin= self.name +': DRIFT, L='+str(self.L)+';'
+ if self.name=='': lin='ELEM'+i+'_DRIFT'+': DRIFT, L='+str(self.L)+';'
+
+ return lin
+
+ def get_fluka(self):
+
+ i='0'*(3-int(log10(self.idx_elem+1)))+str(self.idx_elem+1)
+
+ if self.name!='': lin= self.name +' DRIFT '
+ if self.name=='': lin='ELEM'+i+'_DRIFT'+' DRIFT '
+ lin+=str(self.L)+' '+str(self.s)+' '
+ lin+='0.0 0.0 0.0 '
+ if self.apty==0.0: lin+='CIRCLE '+str(self.apt*1e-3)+' 0.0'
+ if self.apty!=0.0: lin+='RECTANGLE '+str(self.apt*1e-3)+' '+str(self.apty*1e-3)
+
+ return lin
+
+ def get_mars(self):
+
+ lin ='"DRIFT" "" "" '+str(self.s)+' '+str(self.L)+' '
+ lin+='0 0 0 0 0'
+
+ return lin
+
+class BEND(ELEM):
+ '''
+ - For MADX, keeping the sbend+edge structure.
+ - For MADX, if don't export edge, append instances of EDGE.
+ (Otherwise, the matrix is wrong in MADX.)
+ - For MARS/Fluka, assuming the rbend.
+ '''
+ def __init__(self,name,typ,para):
+
+ ELEM.__init__(self,name,typ,para)
+
+ # TW instances
+ self.angle=float(para[0])*pi/180.0 # deg => rad
+ self.rho =float(para[1])*1e-3 # mm => m
+ self.apt =float(para[3])
+ self.plane= para[4]
+
+ # Option instances
+ if self.plane=='0': self.tilt=0.0
+ if self.plane=='1': self.tilt=0.5*pi
+ self.ds =self.rho*abs(self.angle)
+ self.L =2.0*self.rho*sin(0.5*abs(self.angle))
+ self.gap =0.0 # Get this from EDGE IF not exporting edge for MADX
+ self.ff_coef1=0.0 # Get this from EDGE IF not exporting edge for MADX
+ self.ff_coef2=0.0 # Get this from EDGE IF not exporting edge for MADX
+
+ def update_idx(self):
+
+ # Updata idx, idx_elem, s
+ self.idx +=1
+ self.idx_elem+=1
+ self.s +=self.ds # Need an overwrite because of this
+
+ def get_madx(self):
+
+ i='0'*(3-int(log10(self.idx_elem+1)))+str(self.idx_elem+1)
+
+ if self.name!='': lin= self.name +': '
+ if self.name=='': lin='ELEM'+i+'_BEND'+': '
+ lin+='SBEND, L='+str(self.ds )+', ' # sbend
+ #lin+='RBEND, L='+str(self.L )+', ' # rbend
+ lin+='ANGLE=' +str(self.angle )+', '
+ lin+='TILT=' +str(self.tilt )+', '
+ lin+='HGAP=' +str(self.gap*5e-4)+', '
+ lin+='FINT=' +str(self.ff_coef1)+'; '
+
+ return lin
+
+ def get_fluka(self):
+
+ i='0'*(3-int(log10(self.idx_elem+1)))+str(self.idx_elem+1)
+
+ if self.name!='': lin= self.name +' RBEND '
+ if self.name=='': lin='ELEM'+i+'_BEND'+' RBEND '
+ lin+=str(self.ds)+' '+str(self.s)+' '
+ lin+=str(self.tilt)+' '+str(self.angle)+' 0.0 '
+ lin+='CIRCLE '+str(self.apt*1e-3)+' 0.0'
+
+ return lin
+
+ def get_mars(self):
+
+ lin ='"RBEND" "" "" '+str(self.s)+' '+str(self.L)+' '
+ lin+=str(self.angle)+' 0 0 0 '+str(self.tilt)
+
+ return lin
+
+class EDGE(ELEM):
+ '''
+ - For MARS, assuming rbend and the export is commented out.
+ '''
+ def __init__(self,name,typ,para):
+
+ ELEM.__init__(self,name,typ,para)
+
+ # TW instances
+ self.angle =float(para[0])*pi/180.0 # deg => rad
+ self.rho =float(para[1])*1e-3 # mm => m
+ self.gap =float(para[2])
+ self.ff_coef1=float(para[3])
+ self.ff_coef2=float(para[4])
+ self.apt =float(para[5])
+ self.plane = para[6]
+
+ # Option instances
+ if self.plane=='0': self.tilt=0.0
+ if self.plane=='1': self.tilt=0.5*pi
+
+ def get_madx(self):
+
+ i='0'*(3-int(log10(self.idx_elem+1)))+str(self.idx_elem+1)
+
+ if self.name!='': lin= self.name +': DIPEDGE, '
+ if self.name=='': lin='ELEM'+i+'_EDGE'+': DIPEDGE, '
+ lin+='E1=' +str(self.angle) +', ' # sbend
+ #lin+='E1=' +'0' +', ' # rbend
+ lin+='H=' +str(1.0/self.rho) +', '
+ lin+='TILT='+str(self.tilt )+', '
+ lin+='HGAP='+str(self.gap*5e-4)+', '
+ lin+='FINT='+str(self.ff_coef1)+'; '
+
+ return lin
+
+class APERTURE(ELEM):
+ '''
+ - For MADX, tried "collimator" but didn't work ...?
+ '''
+ def __init__(self,name,typ,para):
+
+ ELEM.__init__(self,name,typ,para)
+
+ # TW instances
+ self.apt =float(para[0])
+ self.apty =float(para[1])
+ self.typ_apt= para[2]
+
+ def get_madx(self):
+
+ i='0'*(3-int(log10(self.idx_elem+1)))+str(self.idx_elem+1)
+
+ if self.name!='': lin= self.name +': DRIFT, L=0;'
+ if self.name=='': lin='ELEM'+i+'_APT'+': DRIFT, L=0;'
+
+ ## if self.name!='': lin= self.name +': collimator, l=0, '
+ ## if self.name=='': lin='ELEM'+i+'_apt'+': collimator, l=0, '
+ ## if self.typ_apt=='0':
+ ## lin+='apertype=rectangle, aperture={'+str(self.apt*1e-3)+','+str(self.apty*1e-3)+'};'
+ ## if self.typ_apt=='1':
+ ## lin+='apertype=CIRCLE, aperture={'+str(self.apt*1e-3)+'};'
+
+ return lin
+
+ def get_fluka(self):
+
+ i='0'*(3-int(log10(self.idx_elem+1)))+str(self.idx_elem+1)
+
+ if self.name!='': lin= self.name +' DRIFT '
+ if self.name=='': lin='ELEM'+i+'_APT'+' DRIFT '
+ lin+=str(self.L)+' '+str(self.s)+' '
+ lin+='0.0 0.0 0.0 '
+ if self.typ_apt=='0': lin+='RECTANGLE '+str(self.apt*1e-3)+' '+str(self.apty*1e-3)
+ if self.typ_apt=='1': lin+='CIRCLE '+str(self.apt*1e-3)+' 0.0'
+
+ return lin
+
+ def get_mars(self):
+
+ lin ='"drift" "" "" '+str(self.s)+' '+str(self.L)+' '
+ lin+='0 0 0 0 0'
+
+ return lin
+
+class DIAG_POSITION(ELEM):
+ '''
+ '''
+ def __init__(self,name,typ,para):
+
+ ELEM.__init__(self,name,typ,para)
+
+ def get_madx(self):
+
+ i='0'*(3-int(log10(self.idx_elem+1)))+str(self.idx_elem+1)
+
+ if self.name!='': lin= self.name +': MONITOR, L=0;'
+ if self.name=='': lin='ELEM'+i+'_BPM'+': MONITOR, L=0;'
+
+ return lin
+
+ def get_fluka(self):
+
+ i='0'*(3-int(log10(self.idx_elem+1)))+str(self.idx_elem+1)
+
+ if self.name!='': lin= self.name +' MONITOR '
+ if self.name=='': lin='ELEM'+i+'_BPM'+' MONITOR '
+ lin+=str(self.L)+' '+str(self.s)+' '
+ lin+='0.0 0.0 0.0 '
+ lin+='CIRCLE '+str(self.apt*1e-3)+' 0.0'
+
+ return lin
+
+ def get_mars(self):
+
+ lin ='"MONITOR" "" "" '+str(self.s)+' '+str(self.L)+' '
+ lin+='0 0 0 0 0'
+
+ return lin
+
+class DIAG(ELEM):
+ '''
+ '''
+ def __init__(self,name,typ,para):
+
+ ELEM.__init__(self,name,typ,para)
+
+ def get_madx(self):
+
+ i='0'*(3-int(log10(self.idx_elem+1)))+str(self.idx_elem+1)
+
+ if self.name!='': lin= self.name +': INSTRUMENT, L=0;'
+ if self.name=='': lin='ELEM'+i+'_DIAG'+': INSTRUMENT, L=0;'
+
+ return lin
+
+ def get_fluka(self):
+
+ i='0'*(3-int(log10(self.idx_elem+1)))+str(self.idx_elem+1)
+
+ if self.name!='': lin= self.name +' INSTRUMENT '
+ if self.name=='': lin='ELEM'+i+'_DIAG'+' INSTRUMENT '
+ lin+=str(self.L)+' '+str(self.s)+' '
+ lin+='0.0 0.0 0.0 '
+ lin+='CIRCLE '+str(self.apt*1e-3)+' 0.0'
+
+ return lin
+
+ def get_mars(self):
+
+ lin ='"INSTRUMENT" "" "" '+str(self.s)+' '+str(self.L)+' '
+ lin+='0 0 0 0 0'
+
+ return lin
+
+#---- Active commands
+
+class STEERER(COMM):
+ '''
+ - Only magnetic steerers supported.
+ - Only the ESS type non-linearity supported.
+ '''
+ def __init__(self,name,typ,para):
+
+ COMM.__init__(self,name,typ,para)
+
+ # TW instances
+ self.Bx=float(para[0])
+ self.By=float(para[1])
+
+ # TW option instances
+ try : self.max =float(para[2])
+ except: self.max =0.0
+ try : self.nonlin=float(para[5])
+ except: self.nonlin=0.0
+
+ def get_tw(self):
+
+ if self.name!='': lin=self.name+': '+self.typ+' '
+ if self.name=='': lin= self.typ+' '
+ lin+=str(self.Bx )+' '
+ lin+=str(self.By )+' '
+ lin+=str(self.max )+' '
+ lin+='0' +' '
+ lin+='0' +' '
+ lin+=str(self.nonlin)
+
+ return lin
+
+class CHOPPER(COMM):
+ '''
+ '''
+ def __init__(self,name,typ,para):
+
+ COMM.__init__(self,name,typ,para)
+
+ # TW instances
+ self.Nelem= int(para[0])
+ self.volt =float(para[1])
+ self.gap =float(para[2])*2.0
+ self.d_pos=float(para[3])
+ self.plane= para[4]
+
+ def get_tw(self):
+
+ if self.name!='': lin=self.name+': '+self.typ+' '
+ if self.name=='': lin= self.typ+' '
+ lin+=str(self.Nelem )+' '
+ lin+=str(self.volt )+' '
+ lin+=str(self.gap*0.5)+' '
+ lin+=str(self.d_pos )+' '
+ lin+= self.plane
+
+ return lin
+
+#---- Passive commands
+
+class FREQ(COMM):
+ '''
+ '''
+ def __init__(self,name,typ,para):
+
+ COMM.__init__(self,name,typ,para)
+
+ # TW instances
+ self.freq_new=float(para[0])
+
+ def update_idx(self):
+
+ self.idx +=1
+ self.freq =self.freq_new
+
+class MARKER(COMM):
+ '''
+ '''
+ def __init__(self,name,typ,para):
+
+ COMM.__init__(self,name,typ,para)
+
+ def get_madx(self):
+
+ if self.name!='': lin=self.name +': MARKER;'
+ if self.name=='': lin=self.para[0]+': MARKER;'
+
+ return lin
+
+ def get_fluka(self):
+
+ if self.name!='': lin=self.name +' MARKER '
+ if self.name=='': lin=self.para[0]+' MARKER '
+ lin+='0.0 '+str(self.s)
+
+ return lin
+
+ ## def get_mars(self):
+
+ ## lin ='"marker" "" "" '+str(self.s)+' 0.0 '
+ ## lin+='0 0 0 0 0'
+
+ ## return lin
+
+#---- No longer needed ??
+
+## class ADJUST(COMM):
+## '''
+## '''
+## def __init__(self,name,typ,para):
+
+## COMM.__init__(self,name,typ,para)
+
+## # TW instances
+## self.fam= para[0]
+## self.var=int(para[1])
+
+## # TW option instances
+## try : self.link=self.para[2]
+## except: self.link='0'
+## try : self.min =float(self.para[3])
+## except: self.min =0.0
+## try : self.max =float(self.para[4])
+## except: self.max =0.0
+## try : self.ini =float(self.para[5])
+## except: self.ini =0.0