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

Some further debugging and some nice optics

Browse Source
This commit is contained in:
reiche
2025-05-08 15:08:21 +02:00
parent 3e858452d3
commit 52f2624a4d
9 changed files with 162 additions and 97 deletions
Download Patch File Download Diff File Expand all files Collapse all files

78
model.py
View File

@@ -1,9 +1,15 @@
import json
import copy
import numpy as np
from onlinemodel.core import Facility
from onlinemodel.madx import CMadX
converttwiss= {'betax':'betx','betay':'bety','alphax':'alfx','alphay':'alfy',
'etax':'dx','etay':'dy','etapx':'dpx','etapy':'dpy',
'x':'x','y':'y','px':'px','py':'py','energy':'energy'}
class Model:
def __init__(self, phase=0, parent=None):
print('Initializing online model ...')
@@ -76,7 +82,6 @@ class Model:
undulators={}
kicker={}
energy={'location': self.energyReference, 'energy' : 1e-6*self.om.EnergyAt(self.energyReference)[0]}
for ele in elements:
if 'MQUA' in ele.Name or 'MQSK' in ele.Name:
quadrupoles[ele.Name]={'k1':ele.k1,'k1L':ele.k1*ele.Length}
@@ -94,24 +99,30 @@ class Model:
return {'Quadrupole':quadrupoles,'Sextupole':sextupoles,'Dipole':dipoles,'RF':rf,'Undulator':undulators,
'Kicker':kicker,'Energy':energy, 'InitialCondition':self.startTwiss}
def loadSettingsGroup(self,group,fields):
def loadSettingsGroup(self,group,fields,normalized=False):
for key in group.keys():
ele = self.om.getElement(key)
for field in fields:
ele.__dict__[field]=group[key][field]
if normalized:
ele.__dict__[field]/=ele.Length
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['RF'], ['Gradient'],normalized=True)
self.loadSettingsGroup(settings['RF'], ['Phase'])
self.loadSettingsGroup(settings['Undulator'], ['K','kx','ky'])
self.loadSettingsGroup(settings['Kicker'], ['cory'])
self.startEnergy = settings['Energy']['energy']
self.energyReference = settings['Energy']['location']
self.startTwiss = settings['InitialCondition']
print('initial condition:',self.startTwiss)
print(self.startTwiss)
self.updateEnergy(self.startEnergy)
print(settings['Energy']['energy'])
print('Energy in system',self.om.EnergyAt(self.energyReference)[0])
@@ -130,6 +141,7 @@ class Model:
elif 'K2L' == fld:
ele.k2 = data['settings'][key]/ele.Length
elif 'Grad' == fld:
# ele.Gradient = data['settings'][key]/ele.Length
ele.Gradient = data['settings'][key]
elif 'Phase' == fld:
ele.Phase = data['settings'][key]
@@ -162,10 +174,10 @@ class Model:
self.madx.updateVariables(twiss)
res,twiss,tar=self.madx.match(sequence,variables,conditions,periodic)
energy = self.calcEnergyProfile(twiss)
matchtwiss={'betax':twiss.betx[0],'betay':twiss.betay[0],'alphax':twiss.alfx[0],'alphay':twiss.alfy[0]}
matchtwiss={'betax':twiss.betx[0],'betay':twiss.bety[0],'alphax':twiss.alfx[0],'alphay':twiss.alfy[0]}
if plot:
self.parent.plot.newData(twiss, energy)
return matchtwiss,tar
return res, matchtwiss,tar
def track(self):
@@ -179,7 +191,8 @@ class Model:
refloc = self.parent.reference.getReferenceLocation().upper()
if refloc == 'START':
refloc = start.upper()
twiss0['energy'] = self.startEnergy
twiss0['energy'] = self.startEnergy*1e-3 # convert to GeV for madx
start, end = self.checkRange(start, end, refloc[0:7])
if start is None:
return
@@ -188,13 +201,16 @@ class Model:
def doTrack(self, start,end, refloc, twiss_in, plot = True):
twiss0={key+'0':twiss_in[key] for key in twiss_in.keys()}
twiss0 = {converttwiss[key]:twiss_in[key] for key in twiss_in.keys()}
self.setBranch(end.upper())
if not refloc == start:
twiss0 = self.doBackTrack(refloc,start,twiss0)
self.startTwiss=copy.deepcopy(twiss0)
print('Back Tracking Done')
for key in twiss_in.keys():
if converttwiss[key] in twiss0.keys():
self.startTwiss[key+'0'] = twiss0[converttwiss[key]]
self.madx.updateVariables(twiss0)
twiss = self.madx.track(start+'$START',end+'$END')
twiss = self.madx.track('swissfel',start+'$START',end+'$END')
energy = self.calcEnergyProfile(twiss)
if plot:
self.parent.plot.newData(twiss,energy)
@@ -217,27 +233,29 @@ class Model:
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']
print('Backtracking 1:',twiss0)
twiss0['alfx'] = -twiss0['alfx'] # revert particle trajectories
twiss0['alfy'] = -twiss0['alfy']
twiss0['dpx'] = -twiss0['dpx']
twiss0['dpy'] = -twiss0['dpy']
twiss0['px'] = -twiss0['px']
twiss0['py'] = -twiss0['py']
self.madx.updateVariables(twiss0)
twiss = self.madx.track(start, end+'$START', reverse=True)
print(start,end)
twiss = self.madx.track('invswissfel', start, end+'$START')
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]
twiss1['betx'] = twiss.betx[-1] # revert trajectory back
twiss1['bety'] = twiss.bety[-1]
twiss1['alfx'] = -twiss.alfx[-1]
twiss1['alfy'] = -twiss.alfy[-1]
twiss1['dx'] = twiss.dx[-1]
twiss1['dy'] = twiss.dy[-1]
twiss1['dpx'] = -twiss.dpx[-1]
twiss1['dpy'] = -twiss.dpy[-1]
twiss1['x'] = twiss.x[-1]
twiss1['y'] = twiss.y[-1]
twiss1['px'] = -twiss.px[-1]
twiss1['py'] = -twiss.py[-1]
return twiss1
def setBranch(self,end):
@@ -276,3 +294,5 @@ class Model:
end = ref
print('Reference point is downstream the tracking range. Extending range')
return start,end