SwissFEL-Simulation/OpticsTools
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Basic Tracking done

Browse Source
This commit is contained in:
reiche
2025-05-05 17:52:45 +02:00
parent 93527fc7e5
commit b4c25511d1
10 changed files with 100017 additions and 278 deletions
Download Patch File Download Diff File Expand all files Collapse all files

216
OpticsModel.py
View File

@@ -1,67 +1,170 @@
import sys
import json
import numpy as np
from onlinemodel.core import Facility
from onlinemodel.code import CMadX
from onlinemodel.madx import CMadX
class Model:
def __init__(self, phase=0, parent=None):
print('Initializing online model ...')
self.phase = phase # current planned future
self.parent=parent
self.om = Facility(phase)
self.om = Facility(init=1, alt = phase)
self.order = None
self.madx = CMadX()
self.twiss = None
# hook up events
self.eventHandling()
def getSettings(self):
elements = self.om.listElement('*',1)
quadrupoles={}
sextupoles={}
dipoles={}
rf={}
undulators={}
kicker={}
loc = 'SINLH01.DBAM010'
energy={'location': loc, 'energy':self.om.EnergyAt(loc)}
for ele in elements:
if 'MQUA' in ele.Name:
quadrupoles[ele.Name]={'k1':ele.k1,'k1L':ele.k1*ele.Length}
elif 'MSEX' in ele.Name:
sextupoles[ele.Name]={'k2':ele.k2,'k2L':ele.k2*ele.Length}
elif 'MBND' in ele.Name:
if 'SINLH' in ele.Name or 'SINBC' in ele.Name or 'S10BC' in ele.Name or 'SATMA' in ele.Name or 'SATUN' in ele.Name:
dipoles[ele.Name]={'angle':ele.angle}
elif 'UIND' in ele.Name:
undulators[ele.Name]={'K':ele.K,'kx':ele.kx,'ky':ele.ky}
elif 'RACC' in ele.Name:
rf[ele.Name]={'Gradient':ele.Gradient,'Phase':ele.Phase}
elif 'MKAC' in ele.Name or 'MKDC' in ele.Name:
kicker[ele.Name] = {'cory': ele.cory}
return {'Quadrupole':quadrupoles,'Sextupole':sextupoles,'Dipole':dipoles,'RF':rf,'Undulator':undulators,'Kicker':kicker,'Energy':energy}
def loadSettingsGroup(self,group,fields):
for key in group.keys():
ele = self.om.getElement(key)
for field in fields:
ele.__dict__[field]=group[key][field]
def loadSettings(self,settings):
self.loadSettingsGroup(settings['Quadrupole'],['k1'])
self.loadSettingsGroup(settings['Sextupole'], ['k2'])
self.loadSettingsGroup(settings['Dipole'], ['angle'])
self.loadSettingsGroup(settings['RF'], ['Gradient','Phase'])
self.loadSettingsGroup(settings['Undulator'], ['K','kx','ky'])
self.loadSettingsGroup(settings['Kicker'], ['cory'])
self.om.forceEnergyAt(settings['Energy']['location'],settings['Energy']['energy'][0])
def updateModel(self):
file ='SwissFELRef.json'
with open(file) as f:
data = json.load(f)
for key in data['settings'].keys():
tags=key.split(':')
elename = tags[0].replace('-','.')
fld = tags[1]
ele = self.om.getElement(elename)
if 'K1L' == fld:
ele.k1 = data['settings'][key]/ele.Length
elif 'K2L' == fld:
ele.k2 = data['settings'][key]/ele.Length
elif 'Grad' == fld:
ele.Gradient = data['settings'][key]
elif 'Phase' == fld:
ele.Phase = data['settings'][key]
elif 'kx' == fld:
ele.kx = data['settings'][key]
elif 'ky' == fld:
ele.ky = data['settings'][key]
elif 'K' == fld:
ele.K = data['settings'][key]
elif 'Gap' == fld:
ele.Gap = data['settings'][key]
elif 'Offset' == fld:
ele.Offset = data['settings'][key]
elif 'K0L' == fld:
ele.angle = data['settings'][key]
elif 'KICK' == fld:
ele.cory=data['settings'][key]
else:
print('Unknown tag:', key, data['settings'][key])
def eventHandling(self):
self.parent.UITrack.clicked.connect(self.track)
def track(self):
start = str(self.parent.UITrackStart.text())
end = str(self.parent.UITrackEnd.text())
twiss = {}
twiss['betax0'] = float(str(self.parent.UIBetax.text()))
twiss['betay0'] = float(str(self.parent.UIBetay.text()))
twiss['alphax0'] = float(str(self.parent.UIAlphax.text()))
twiss['alphay0'] = float(str(self.parent.UIAlphay.text()))
twiss['energy0'] = 150
self.doTrack(start,end,'SINLH02.MQUA410$START',twiss)
start = str(self.parent.UITrackStart.text()).upper()
end = str(self.parent.UITrackEnd.text()).upper()
if len(start)>7:
start = start[0:7]
if len(end)>7:
end = end[0:7]
twiss0 = self.parent.reference.getReferencePoint()
refloc = self.parent.reference.getReferenceLocation().upper()
twiss0['energy'] = 150.
start, end = self.checkRange(start, end, refloc[0:7])
if start is None:
return
print('Tracking from',start,'to',end)
self.doTrack(start,end,refloc,twiss0)
def doBackTrack(self,start=None,end=None,twiss=None):
twiss['alphax0'] = -twiss['alphax0']
twiss['alphay0'] = -twiss['alphay0']
self.madx.updateVariables(twiss)
self.madx.updateLattice(self.om)
twiss = self.madx.track(start, end, reverse=True)
#self.parent.plot.newData(twiss)
twiss0 = {}
twiss0['betax0']= twiss.betx[-1]
twiss0['betay0'] = twiss.bety[-1]
twiss0['alphax0'] = -twiss.alfx[-1]
twiss0['alphay0'] = -twiss.alfy[-1]
self.madx.clear()
return twiss0
def doTrack(self, start,end,refloc, twiss):
self.setBranch(start,end)
# check if refloc is in the range of start and end
def doTrack(self, start,end, refloc, twiss_in, plot = True):
twiss0={key+'0':twiss_in[key] for key in twiss_in.keys()}
self.setBranch(end.upper())
if not refloc == start:
twiss = self.doBackTrack(refloc,start,twiss)
twiss0 = self.doBackTrack(refloc,start,twiss0)
self.madx.updateVariables(twiss0)
twiss = self.madx.track(start+'$START',end+'$END')
self.madx.updateVariables(twiss)
self.madx.updateLattice(self.om)
energy = np.array([0. for i in range(len(twiss.betx))])
twiss = self.madx.track()
self.parent.plot.newData(twiss)
e0 = 0.
de = 0
for i, name in enumerate(twiss.name):
if len(name) > 15:
elename = name[0:15]
ele = self.om.getElement(elename)
if not ele is None:
erg= self.om.EnergyAt(ele)
if e0 == 0.:
energy[:i]+=erg[0]*1e-6
e0 = (erg[0]+erg[1])*1e-6
energy[i]=e0
if plot:
self.parent.plot.newData(twiss,energy)
def setBranch(self,start,end):
start0 = start[0:7]
end0 = end[0:7]
def doBackTrack(self,start=None,end=None,twiss0=None):
twiss0['alphax0'] = -twiss0['alphax0'] # revert particle trajectories
twiss0['alphay0'] = -twiss0['alphay0']
twiss0['etapx0'] = -twiss0['etapx0']
twiss0['etapy0'] = -twiss0['etapx0']
twiss0['px0'] = -twiss0['px0']
twiss0['py0'] = -twiss0['py0']
self.madx.updateVariables(twiss0)
twiss = self.madx.track(start, end+'$START', reverse=True)
twiss1 = {}
twiss1['betax0'] = twiss.betx[-1] # revert trajectory back
twiss1['betay0'] = twiss.bety[-1]
twiss1['alphax0'] = -twiss.alfx[-1]
twiss1['alphay0'] = -twiss.alfy[-1]
twiss1['etax0'] = twiss.dx[-1]
twiss1['etay0'] = twiss.dy[-1]
twiss1['etapx0'] = -twiss.dpx[-1]
twiss1['etapy0'] = -twiss.dpy[-1]
twiss1['x0'] = twiss.x[-1]
twiss1['y0'] = twiss.y[-1]
twiss1['px0'] = -twiss.px[-1]
twiss1['py0'] = -twiss.py[-1]
return twiss1
def setBranch(self,end):
destination = 'ARAMIS'
if 'SPO' in end:
destination = 'PORTHOS'
@@ -69,4 +172,31 @@ class Model:
destination = 'ATHOS'
elif 'S10BD' in end or 'SIN' in end:
destination = 'INJECTOR'
self.om.setBranch(destination,start0,end0)
self.om.setBranch(destination,'SINLH01')
self.order=self.om.getBranchElements()
self.madx.updateLattice(self.om,destination)
def checkRange(self,start,end,ref):
if self.om.getSection(start) is None:
print('Invalid staring point for tracking. Setting to SINLH01')
startsec='SINLH01'
else:
startsec = self.om.getSection(start).Element[0].Name
if self.om.getSection(end) is None:
print('Invalid staring point for tracking. Setting to starting section')
endsec = startsec
else:
endsec = self.om.getSection(end).Element[0].Name
if self.om.getSection(ref) is None:
print('Invalid section for reference. Aborting tracking')
return None,None
refsec = self.om.getSection(ref).Element[0].Name
if not self.om.isUpstream(startsec,endsec):
return self.checkRange(end,start,ref)
if not self.om.isUpstream(startsec,refsec):
start = ref
print('Reference point is upstream the tracking range. Extending range')
if not self.om.isUpstream(refsec,endsec):
end = ref
print('Reference point is downstream the tracking range. Extending range')
return start,end