SwissFEL-Simulation/ElegantTools
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
dev
ElegantTools/lattice.py
T
reiche 56a4aff6ac First release of ElegantTools
2026-01-26 11:35:00 +01:00

372 lines
16 KiB
Python
Raw Permalink Blame History

import math
from onlinemodel.core.apptemplate import ApplicationTemplate
# Elegant Interface
class Lattice(ApplicationTemplate):
def __init__(self,sessionID='.'):
ApplicationTemplate.__init__(self)
self.wpath='ReferenceFiles' # path for reference files
self.fid = 0
self.lsc = 0 # flag whether lsc is enabled
self.csr = 0 # flag whether csr is enabled or not
self.wakebins=0 # automatic calculation of
self.charge = 200e-12 # default charge in the lattice
self.dump = [] # diagnostics location where to dump particle distribution
self.seq=[]
self.driftcount = 1 # drift counter
self.lscbins=100 # number of bins for lsc calculation
self.csrbins=100 # number of bins for CSR calculation
self.csrmaxlen=10. # distance over which the CSR wake is considered after a dipole
self.csrlen=0 # running variable to count distance after last bend
self.center=[] # location where the beam is centered due to CSR kicks
self.Distribution = 1
def openLatticeStream(self,path,name):
filename='%s/%s.lat' % (path,name)
# self.path=path
self.path='.'
self.savepath=path
self.seq=[]
self.fid=open(filename,'w')
self.driftcount=1
self.write('Q: CHARGE, TOTAL=%e\n' % self.charge)
self.write('ALTER: CENTER, DELTA=1,T=1\n') # put the beam on the right energy
return filename
def closeLatticeStream(self):
self.fid.close()
# routines to write lattice
def setCSR(self):
self.csrlen=self.csrmaxlen
def isType(self,name):
if (name.find('elegant')>-1):
return 1
else:
return 0
def write(self,line):
self.fid.write(line)
def writeLine(self,line,seq):
lc=1
name=line.Name
seqloop=seq
if len(name)==7:
seqloop=self.seq
if len(name) < 1:
name='SwissFEL'
self.write("\n%s.Start: MARKER\n" % (name))
secend = name+'.End'
if secend in self.dump:
print('Adding dump in Elegant output for',name)
self.write('%s.End: WATCH, FILENAME="Runs/%%s.%s.End.out", MODE=COORD\n' % (name,name))
else:
self.write("%s.End: MARKER\n" % (name))
if (name=="SwissFEL"):
self.write("\n%s: Line=(Q, ALTER, %s.Start, " % (name,name))
else:
self.write("\n%s: Line=(%s.Start, " % (name,name))
for ele in seqloop:
# print ele
if isinstance(ele,str):
self.write("%s, " % (ele))
else:
self.write("%s, " % (ele['Name']))
lc=lc+1
if (lc>5):
self.write("&\n\t")
lc=0
self.write("%s.End)\n\n" % (name))
if len(name)==7:
del self.seq[:]
def writeDrift(self,l):
if l <= 0 :
return 0
name="Drift%4.4d" % (self.driftcount)
if self.csrlen > 0:
self.write("%s.csr: CSRDRIFT, L=%f, DZ=0.1, USE_STUPAKOV=1, CSR=%d\n" % (name,l,self.csr))
self.write("%s.lsc: LSCDRIFT, L=0, LEFFECTIVE=%f, BINS=%d, SMOOTHING=1, LSC=%d\n" % (name,l,self.lscbins,self.lsc))
self.write("%s: LINE=(%s.csr,%s.lsc)\n" % (name,name,name))
self.csrlen=self.csrlen-l
else:
self.write("%s: LSCDRIFT, L=%f, LEFFECTIVE=%f, BINS=%d, SMOOTHING=1, LSC=%d\n" % (name,l,l,self.lscbins,self.lsc))
self.seq.append(name)
self.driftcount=self.driftcount+1
return name
def writeMarker(self,ele):
if 'ALIG' in ele.Name and 'dy' in ele.__dict__.keys():
self.write('%s: MALIGN, DY=-%f \n' % (ele.Name,ele.dy))
else:
self.write('%s: MARK \n' % (ele.Name))
self.seq.append(ele.Name)
def writeVacuum(self,ele):
ds=ele.sRef+0.5*ele.getLength();
if ds>0:
self.writeDrift(ds)
self.writeMarker(ele)
ds=ele.getResLength()-ele.sRef-0.5*ele.getLength();
if ds>0:
self.writeDrift(ds)
def writeDiagnostic(self,ele):
Seval=ele.Seval
if (Seval<0):
Seval=0.5*ele.Length
ds=ele.sRef+Seval;
if ds>0:
self.writeDrift(ds)
if ele.Name in self.center:
print('Adding center in Elegant output for',ele.Name)
self.write('%s_CENTER: CENTER\n' % ele.Name)
self.seq.append('%s_CENTER' % ele.Name)
if ele.Name in self.dump:
print('Adding dump in Elegant output for',ele.Name)
self.write('%s: WATCH, FILENAME="Runs/%%s.%s.out", MODE=COORD\n' % (ele.Name,ele.Name))
self.seq.append(ele.Name)
else:
self.writeMarker(ele)
ds=ele.getResLength()-ds
if ds>0:
self.writeDrift(ds)
def writeCorrector(self,ele):
corx=0
cory=0
if 'corx' in ele.__dict__:
corx=ele.corx
if 'cory' in ele.__dict__:
cory=ele.cory
if 'MADthin' in ele.__dict__:
if 'MADshift' in ele.__dict__:
ds=ele.Length*0.5+ele.MADshift
self.writeDrift(ds)
self.write("%s: KICKER, L=%f, HKICK=%e,VKICK=%e\n" % (ele.Name,0,corx,cory))
self.seq.append(ele.Name)
return
if (ele.sRef>0):
self.writeDrift(ele.sRef)
if self.lsc==0 or ele.getLength()==0:
self.write("%s: KICKER, L=%f, HKICK=%e,VKICK=%e\n" % (ele.Name,ele.getLength(),corx,cory))
self.seq.append(ele.Name)
else:
self.write("%s.cor: KICKER, L=%f, HKICK=%e,VKICK=%e\n" % (ele.Name,ele.getLength(),corx,cory))
self.write("%s.lsc: LSCDRIFT, L=0, LEFFECTIVE=%f, BINS=%d, SMOOTHING=1, LSC=%d\n" % (ele.Name,ele.getLength(),self.lscbins,self.lsc))
self.write("%s: LINE=(%s.cor,%s.lsc)\n" % (ele.Name,ele.Name,ele.Name))
self.seq.append(ele.Name)
ds=ele.getResLength()-ele.getLength()-ele.sRef
if ds > 0:
self.writeDrift(ds)
def writeSextupole(self,ele):
if (ele.sRef>0):
self.writeDrift(ele.sRef)
if self.lsc==0:
self.write("%s: SEXTUPOLE, L=%f, K2=%e, Tilt=%f\n" % (ele.Name,ele.getLength(),ele.k2,ele.Tilt))
self.seq.append(ele.Name)
else:
self.write("%s.sex: SEXTUPOLE, L=%f, K2=%e, Tilt=%f\n" % (ele.Name,ele.getLength(),ele.k2,ele.Tilt))
self.write("%s.lsc: LSCDRIFT, L=0, LEFFECTIVE=%f, BINS=%d, SMOOTHING=1, LSC=%d\n" % (ele.Name,ele.getLength(),self.lscbins,self.lsc))
self.fid.write("%s: LINE=(%s.sex,%s.lsc)\n" % (ele.Name,ele.Name,ele.Name))
self.seq.append(ele.Name)
ds=ele.getResLength()-ele.getLength()-ele.sRef
if ds > 0:
self.writeDrift(ds)
def writeQuadrupole(self,ele):
if (ele.sRef>0):
self.writeDrift(ele.sRef)
if self.lsc!=0 or ('corx' in ele.__dict__) or 'cory' in ele.__dict__:
locseq=[]
if self.lsc>0:
self.write("%s.lsc: LSCDRIFT, L=0, LEFFECTIVE=%f, BINS=%d, SMOOTHING=1, LSC=%d\n" % (ele.Name,ele.getLength(),self.lscbins,self.lsc))
locseq.append("%s.lsc" % (ele.Name))
if ('corx' in ele.__dict__) or 'cory' in ele.__dict__:
corx=0
cory=0
if 'corx' in ele.__dict__:
corx=ele.corx
if 'cory' in ele.__dict__:
corx=ele.cory
self.write("%s.q1: QUADRUPOLE, L=%f, K1=%e, Tilt=%f\n" % (ele.Name,ele.getLength()*0.5,ele.k1,ele.Tilt))
self.write("%s.cor: KICKER, L=0, HKICK=%e,VKICK=%e\n" % (ele.Name,corx,cory))
self.write("%s.q2: QUADRUPOLE, L=%f, K1=%e, Tilt=%f\n" % (ele.Name,ele.getLength()*0.5,ele.k1,ele.Tilt))
locseq.append("%s.q1" % (ele.Name))
locseq.append("%s.cor" % (ele.Name))
locseq.append("%s.q2" % (ele.Name))
else:
self.write("%s.q: QUADRUPOLE, L=%f, K1=%e, Tilt=%f\n" % (ele.Name,ele.getLength(),ele.k1,ele.Tilt))
locseq.append("%s.q" % (ele.Name))
self.fid.write("%s: LINE=(%s" % (ele.Name,locseq[0]))
for p in locseq[1:len(locseq)]:
self.fid.write(",%s" % p)
self.fid.write(")\n")
self.seq.append(ele.Name)
else:
self.write("%s: QUADRUPOLE, L=%f, K1=%e, Tilt=%f\n" % (ele.Name,ele.getLength(),ele.k1,ele.Tilt))
self.seq.append(ele.Name)
ds=ele.getResLength()-ele.getLength()-ele.sRef
if ds > 0:
self.writeDrift(ds)
def writeBend(self,ele):
if (ele.sRef>0):
self.writeDrift(ele.sRef)
Lpath=ele.getLength()
angrad=ele.angle*math.asin(1)/90
angradcor=angrad
if 'cor' in ele.__dict__:
angradcor=angrad+ele.cor
if (self.lsc!=0):
self.write("%s.lsc: LSCDRIFT, L=0, LEFFECTIVE=%f, BINS=%d, SMOOTHING=1, LSC=%d\n" % (ele.Name,Lpath,self.lscbins,self.lsc))
self.seq.append("%s.lsc" % ele.Name)
self.write("%s.b: CSRCSBEND, L=%f, ANGLE=%f, &\n" % (ele.Name,Lpath,angradcor))
self.write("\tE1=%f, E2=%f, TILT=%f, &\n" % (ele.e1*angrad,ele.e2*angrad,ele.Tilt))
self.write("\tHGAP=0.015, FINT=0.5, NONLINEAR=1, N_KICKS=10,INTEGRATION_ORDER=4, &\n")
self.write("\tBINS=%d,SG_HALFWIDTH=2, ISR=%d, SYNCH_RAD=%d, CSR=%d\n" % (self.csrbins,self.csr,self.csr,self.csr))
self.write("%s: LINE=(%s.lsc,%s.b)\n" % (ele.Name,ele.Name,ele.Name))
self.seq.append(ele.Name)
else:
self.write("%s: CSRCSBEND, L=%f, ANGLE=%f, &\n" % (ele.Name,Lpath,angradcor))
self.write("\tE1=%f, E2=%f, TILT=%f, &\n" % (ele.e1*angrad,ele.e2*angrad,ele.Tilt))
self.write("\tHGAP=0.015, FINT=0.5, NONLINEAR=1, N_KICKS=10,INTEGRATION_ORDER=4, &\n")
self.write("\tBINS=%d,SG_HALFWIDTH=2, ISR=%d, SYNCH_RAD=%d, CSR=%d\n" % (self.csrbins,self.csr,self.csr,self.csr))
self.seq.append(ele.Name)
if abs(angrad)>0:
self.setCSR()
ds=ele.getResLength()-ele.getLength()-ele.sRef
if ds > 0:
self.writeDrift(ds)
def writeUndulator(self,ele):
pi=2*math.asin(1)
Nu=round(ele.getLength()*ele.ku/pi)
r=3e8/511000/2/pi/100;
Bu=ele.K*ele.ku*0.01/2/pi/r;
LaserHeater=0
if 'UPHS' in ele.Name:
ele.K=0
if 'UDLY' in ele.Name:
ele.K=0
if 'Power' in ele.__dict__:
if ele.Power>0:
LaserHeater=0 ######## disable laser modulator, since energy spread is not correct anyhow
if ele.sRef>0:
self.writeDrift(ele.sRef)
if ele.K==0:
self.write("%s: LSCDRIFT, L=%f, LEFFECTIVE=%f, BINS=%d, SMOOTHING=1, LSC=%d\n" % (ele.Name,ele.getLength(),ele.getLength(),self.lscbins,self.lsc))
self.seq.append(ele.Name)
# add dechirper here
#if 'UDCP' in ele.Name:
# wake=Dechirper()
# wake.getWake(ele.gap, ele.Name,self.savepath)
# print('Write wakefile for %s with gap: %f mm' % (ele.Name,ele.gap*1e3))
# self.write('%s_LW : WAKE, INPUTFILE = "%s/wake_%s.sdds", CHANGE_P0 = 1, &\n' % (ele.Name,self.path,ele.Name))
# self.write('\t TCOLUMN="t",WCOLUMN="W", INTERPOLATE=1, FACTOR = %f, N_BINS=%d, SMOOTHING=1 \n'% (ele.getLength(),self.wakebins))
# self.seq.append('%s_LW' % ele.Name)
# if 'DECHIRPER-V' in ele.Baugruppe:
# self.write('%s_DW : TRWAKE, INPUTFILE = "%s/wake_%s.sdds", &\n' % (ele.Name,self.path,ele.Name))
# self.write('\t TCOLUMN="t",WXCOLUMN="WD", WYCOLUMN="WD", XFACTOR = 0, YFACTOR = %f, &\n' % (ele.getLength()))
# self.write('\t INTERPOLATE=1, N_BINS=%d, SMOOTHING=1 \n'% self.wakebins)
# self.write('%s_QW : TRWAKE, INPUTFILE = "%s/wake_%s.sdds", &\n' % (ele.Name,self.path,ele.Name))
# self.write('\t TCOLUMN="t",WXCOLUMN="WQP", WYCOLUMN="WQ", XFACTOR = %f, YFACTOR = %f, &\n' % (ele.getLength(),ele.getLength()))
# self.write('\t INTERPOLATE=1, N_BINS=%d, SMOOTHING=1, & \n'% self.wakebins)
# self.write('\t X_DRIVE_EXPONENT=0, Y_DRIVE_EXPONENT = 0, X_PROBE_EXPONENT = 1, Y_PROBE_EXPONENT = 1\n')
# self.seq.append('%s_DW' % ele.Name)
# self.seq.append('%s_QW' % ele.Name)
# elif 'DECHIRPER-H' in ele.Baugruppe:
# self.write('%s_DW : TRWAKE, INPUTFILE = "%s/wake_%s.sdds", &\n' % (ele.Name,self.path,ele.Name))
# self.write('\t TCOLUMN="t",WXCOLUMN="WD", WYCOLUMN="WD", XFACTOR = %f, YFACTOR = 0, &\n' % (ele.getLength()))
# self.write('\t INTERPOLATE=1, N_BINS=%d, SMOOTHING=1 \n'% self.wakebins)
# self.write('%s_QW : TRWAKE, INPUTFILE = "%s/wake_%s.sdds", &\n' % (ele.Name,self.path,ele.Name))
# self.write('\t TCOLUMN="t",WXCOLUMN="WQ", WYCOLUMN="WQP", XFACTOR = %f, YFACTOR = %f, &\n' % (ele.getLength(),ele.getLength()))
# self.write('\t INTERPOLATE=1, N_BINS=%d, SMOOTHING=1, & \n'% self.wakebins)
# self.write('\t X_DRIVE_EXPONENT=0, Y_DRIVE_EXPONENT = 0, X_PROBE_EXPONENT = 1, Y_PROBE_EXPONENT = 1\n')
# self.seq.append('%s_DW' % ele.Name)
# self.seq.append('%s_QW' % ele.Name)
ds=ele.getResLength()-ele.getLength()-ele.sRef
if ds>0:
self.writeDrift(ds)
return
name=ele.Name
if self.lsc!=0:
self.write("%s.lsc: LSCDRIFT, L=0, LEFFECTIVE=%f, BINS=%d, SMOOTHING=1, LSC=%d\n" % (ele.Name,ele.getLength(),self.lscbins,self.lsc))
self.seq.append("%s.lsc" % ele.Name)
name="%s.wig" % ele.Name
if LaserHeater==0:
self.write("%s : WIGGLER, L=%f, K=%f, POLES=%d, TILT=%f\n" % (name,ele.getLength(),ele.K,Nu,ele.Tilt))
else:
self.write("%s : LSRMDLTR, L=%f, PERIODS=%d, BU=%f, &\n" % (name, ele.getLength(),Nu/2,Bu));
self.write("\t LASER_W0=%e, LASER_PEAK_POWER=%e, &\n" %(ele.Waist,ele.Power));
self.write("\t N_STEPS=%d, ACCURACY=0.001, POLE_FACTOR1=1, POLE_FACTOR2=1,POLE_FACTOR3=1\n" % (Nu*60));
self.seq.append(name)
ds=ele.getResLength()-ele.getLength()-ele.sRef
if ds>0:
self.writeDrift(ds)
def writeRF(self,ele):
if ele.sRef>0:
self.writeDrift(ele.sRef)
if ele.Tag=="RACC":
self.write('%s : RFCW, FREQ=%f, CELL_LENGTH=1.0, &\n' % (ele.Name,ele.Frequency))
self.write('\t L=%f, VOLT=%e, PHASE= %f, &\n' % (ele.getLength(),ele.Gradient*ele.getLength(),ele.Phase))
self.write('\t CHANGE_P0=1, END1_FOCUS=1, END2_FOCUS=1, &\n');
self.write('\t BODY_FOCUS_MODEL="NONE",INTERPOLATE=1, N_BINS=%d, SMOOTHING=1, &\n' % (self.wakebins));
self.write('\t ZWAKEFILE="%s/wake_l_%s-band.sdds", &\n' % (self.wpath,ele.Band));
self.write('\t TRWAKEFILE="%s/wake_t_%s-band.sdds", &\n' % (self.wpath,ele.Band));
self.write('\t LSC=%d, LSC_BINS=%d, N_KICKS=10, &\n' % (self.lsc,self.lscbins));
self.write('\t TCOLUMN="t",WXCOLUMN="W", WYCOLUMN="W", WZCOLUMN="W"\n');
else:
if ele.Band=='S':
SWave=1
else:
SWave=0
self.write('%s : RFDF, FREQUENCY=%f, STANDING_WAVE=%d, TILT=%f, &\n' % (ele.Name,ele.Frequency,SWave,ele.Tilt))
grad= ele.Gradient*ele.getLength()
self.write('\t L=%f, VOLTAGE=%e, PHASE= %f\n' % (ele.getLength(),grad,ele.Phase))
self.seq.append(ele.Name)
ds=ele.getResLength()-ele.getLength()-ele.sRef
if ds>0:
self.writeDrift(ds)
Reference in New Issue View Git Blame Copy Permalink