424 lines
16 KiB
Python
424 lines
16 KiB
Python
import json
|
|
import copy
|
|
import numpy as np
|
|
|
|
from onlinemodel.core import Facility
|
|
from onlinemodel.madx import CMadX
|
|
#from sipbuild.generator.outputs.formatters import variable
|
|
|
|
converttwiss= {'betax':'betx','betay':'bety','alphax':'alfx','alphay':'alfy',
|
|
'etax':'dx','etay':'dy','etapx':'dpx','etapy':'dpy','mux':'mux','muy':'muy',
|
|
'x':'x','y':'y','px':'px','py':'py','energy':'energy'}
|
|
|
|
|
|
maxgradient= {'QFF':50.,'QFM':50,'QFD':20,'QFS':20,'QFDM':20,'QFA':31.75,'HFA':37.,'HFB':1445}
|
|
|
|
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(init=1, alt = phase)
|
|
self.order = None
|
|
|
|
self.madx = CMadX()
|
|
self.startTwiss = None
|
|
self.startEnergy = None
|
|
self.energyReference ='SINLH02.MBND100'
|
|
|
|
# flag to enfore new lattice
|
|
self.forceLat=True
|
|
# hook up events
|
|
self.matchplot=True
|
|
self.eventHandling()
|
|
|
|
def eventHandling(self):
|
|
self.parent.UITrack.clicked.connect(self.track)
|
|
#self.parent.UIMatchSelected.clicked.connect(self.match)
|
|
|
|
def forceLatticeUpdate(self):
|
|
self.forceLat=True
|
|
|
|
def getLatticeVersion(self):
|
|
return self.om.Version
|
|
|
|
def getInitialEnergy(self):
|
|
return self.om.EnergyAt(self.energyReference)[0]
|
|
|
|
def getVariableInfo(self,varlist):
|
|
variable = {}
|
|
for var in varlist:
|
|
if isinstance(var,dict):
|
|
for key in var.keys():
|
|
variable[key] = {'Val':var[key],'Max':None}
|
|
continue
|
|
ele = self.om.getElement(var.replace('-','.'))
|
|
if ele is None:
|
|
variable[var]={'Val':0,'Max':None}
|
|
else:
|
|
if 'MQUA' in var or 'MQSK' in var or 'MSEX' in var:
|
|
bg = ele.Baugruppe
|
|
Brho = self.om.EnergyAt(var.replace('-','.'))[0]/3e8
|
|
if bg in maxgradient.keys():
|
|
mgrad=maxgradient[bg]/Brho
|
|
else:
|
|
mgrad=None
|
|
if 'MSEX' in var:
|
|
variable[var]={'Val':ele.k2,'Max':mgrad}
|
|
else:
|
|
variable[var] = {'Val': ele.k1, 'Max': mgrad}
|
|
else:
|
|
variable[var]={'Val':0,'Max':None}
|
|
return variable
|
|
|
|
|
|
def updateFromMachine(self,machine):
|
|
pol = ['LH', 'LV+', 'LV-', 'C+', 'C-', 'ZL']
|
|
mag = machine['Magnet']
|
|
for key in mag:
|
|
keyom=key.replace('-','.')
|
|
self.updateElement(keyom,[mag[key]]) # needs to be a list
|
|
und = machine['Undulator']
|
|
for key in und:
|
|
keyom=key.replace('-','.')
|
|
self.updateElement(keyom,[und[key][0],pol[int(und[key][1])]])
|
|
rf = machine['RF']
|
|
for key in rf:
|
|
keyom=key.replace('-','.')
|
|
self.updateElement(keyom,rf[key])
|
|
energy = machine['Energy']
|
|
for key in energy:
|
|
self.updateEnergy(energy[key])
|
|
self.forceLatticeUpdate()
|
|
|
|
def updateEnergy(self,E0):
|
|
if isinstance(E0,list):
|
|
E0=E0[0]
|
|
self.om.forceEnergyAt(self.energyReference, E0*1e6)
|
|
|
|
def updateElement(self,name,val):
|
|
if 'MQUA' in name or 'MQSK' in name:
|
|
L = self.om.getRegExpElement(name[0:7], name[8:15],'Length')[0]
|
|
self.om.setRegExpElement(name[0:7], name[8:15], 'k1', float(val[0])/L)
|
|
if 'MSEX' in name:
|
|
L = self.om.getRegExpElement(name[0:7], name[8:15], 'Length')[0]
|
|
self.om.setRegExpElement(name[0:7], name[8:15], 'k2', float(val[0])/L)
|
|
if 'MBND' in name:
|
|
self.om.setRegExpElement(name[0:7], 'MBND', 'angle', float(val[0]))
|
|
if 'UMOD' in name:
|
|
self.om.setRegExpElement(name[0:7], 'UMOD', 'K', float(val[0]))
|
|
if 'UIND' in name:
|
|
self.om.setRegExpElement(name[0:7], name[8:15], 'K', float(val[0]))
|
|
if 'SATUN' in name:
|
|
kx = 0
|
|
ky = 1
|
|
if 'LV' in val[1]:
|
|
kx = 1
|
|
ky = 0
|
|
if 'C' in val[1]:
|
|
kx = 0.5
|
|
ky=0.5
|
|
self.om.setRegExpElement(name[0:7], name[8:15], 'kx', kx)
|
|
self.om.setRegExpElement(name[0:7], name[8:15], 'ky', ky)
|
|
if 'RSYS' in name:
|
|
grad = float(val[0])
|
|
phase = float(val[1])
|
|
if 'CB' in name[0:7]:
|
|
grad = grad/ 4.
|
|
elif 'XB' in name[0:7] or 'SINSB03' in name or 'SINSB04' in name:
|
|
grad = grad/2
|
|
L = self.om.getRegExpElement(name[0:7], 'RACC', 'Length')[0]
|
|
self.om.setRegExpElement(name[0:7], 'RACC', 'Gradient', grad/L)
|
|
self.om.setRegExpElement(name[0:7], 'RACC', 'Phase', phase)
|
|
|
|
|
|
def getElements(self):
|
|
return self.om.listElement('*', 1)
|
|
|
|
def getSettings(self):
|
|
elements = self.getElements()
|
|
quadrupoles={}
|
|
sextupoles={}
|
|
dipoles={}
|
|
rf={}
|
|
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}
|
|
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 or 'UMOD' 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*ele.Length,'Phase':ele.Phase}
|
|
elif 'MKAC' in ele.Name or 'MKDC' in ele.Name:
|
|
kicker[ele.Name] = {'cory': ele.cory,'design_kick':ele.design_kick}
|
|
return {'Quadrupole':quadrupoles,'Sextupole':sextupoles,'Dipole':dipoles,'RF':rf,'Undulator':undulators,
|
|
'Kicker':kicker,'Energy':energy, 'InitialCondition':self.startTwiss}
|
|
|
|
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'],normalized=True)
|
|
self.loadSettingsGroup(settings['RF'], ['Phase'])
|
|
self.loadSettingsGroup(settings['Undulator'], ['K','kx','ky'])
|
|
self.loadSettingsGroup(settings['Kicker'], ['cory','design_kick'])
|
|
self.startEnergy = settings['Energy']['energy']
|
|
self.energyReference = settings['Energy']['location']
|
|
self.startTwiss = settings['InitialCondition']
|
|
print(self.startTwiss)
|
|
self.updateEnergy(self.startEnergy)
|
|
print(settings['Energy']['energy'])
|
|
print('Energy in system',self.om.EnergyAt(self.energyReference)[0])
|
|
self.forceLat=True
|
|
|
|
|
|
##############################33
|
|
# tracking
|
|
|
|
|
|
def match(self):
|
|
config = self.parent.reference.getMatchingPoint()
|
|
if config is False:
|
|
return
|
|
self.setBranch(config['destination'])
|
|
sequence = config['sequence']
|
|
ID = config['ID']
|
|
start=config['start']['Location']
|
|
end=config['end']
|
|
if len(end) == 7:
|
|
end=end+'$end'
|
|
if 'Twiss' in config['start']:
|
|
itwiss = {converttwiss[key.lower()]: config['start']['Twiss'][key] for key in
|
|
config['start']['Twiss'].keys()}
|
|
else:
|
|
itwiss = None
|
|
|
|
|
|
var={}
|
|
nvar=0
|
|
for key in config['variable'].keys():
|
|
nvar+=1
|
|
key0=key.replace('-','.').lower()
|
|
if 'mqua' in key0 or 'mqsk' in key0:
|
|
key0+='.k1'
|
|
elif 'msex' in key0:
|
|
key0+='.k2'
|
|
elif 'mkac' in key0 or 'mkdc' in key0:
|
|
key0+='.cory'
|
|
var[key0]=config['variable'][key]
|
|
|
|
target = config['target']
|
|
if 'Preset' in target:
|
|
self.updateModelFromMatching(target['Preset'])
|
|
self.parent.reference.updateMatchPoint(ID, 0)
|
|
return
|
|
|
|
if 'Script' in target.keys():
|
|
print(var)
|
|
res,twiss,err = self.madx.callScript(script = target['Script'],sequence=sequence, start=start, end=end,init=itwiss,var=var,preset=True)
|
|
self.updateModelFromMatching(res)
|
|
self.parent.reference.updateMatchPoint(ID, err)
|
|
if self.matchplot:
|
|
energy = self.calcEnergyProfile(twiss)
|
|
self.parent.plot.newData(twiss, energy)
|
|
return
|
|
|
|
ncon = 0
|
|
condilist = {}
|
|
for key in target.keys():
|
|
condi={}
|
|
for ele in target[key]:
|
|
if isinstance(ele,tuple):
|
|
condi[converttwiss[ele[0].lower()]]={'Condition': ele[1],'Val':ele[2]}
|
|
ncon+=1
|
|
elif isinstance(ele,str):
|
|
print('Needs matching results from', ele)
|
|
return
|
|
condilist[key]=condi
|
|
print('Variables:',nvar)
|
|
print('Conditions:',ncon)
|
|
if nvar > ncon:
|
|
print('Adding dummy constraints from initial conditions')
|
|
if not start in condilist.keys():
|
|
condilist[start]={}
|
|
condi={}
|
|
for twkey in itwiss.keys():
|
|
if ncon < nvar:
|
|
condi[twkey]={'Condition':0,'Val':itwiss[twkey]}
|
|
ncon+=1
|
|
condilist[start]=condi
|
|
|
|
for key in condilist.keys():
|
|
print(key,condilist[key])
|
|
|
|
random = self.parent.UIMatchRandom.isChecked()
|
|
res,twiss,err=self.madx.match(sequence=sequence, start=start, end=end,
|
|
init=itwiss, var=var, const=condilist,
|
|
preset=False,random=random)
|
|
|
|
self.updateModelFromMatching(res)
|
|
if config['save'] == True:
|
|
self.parent.reference.saveTwiss(ID,twiss)
|
|
self.parent.reference.updateMatchPoint(ID,err)
|
|
if self.matchplot:
|
|
energy = self.calcEnergyProfile(twiss)
|
|
self.parent.plot.newData(twiss,energy)
|
|
|
|
|
|
def updateModelFromMatching(self,var):
|
|
for magm in var.keys():
|
|
mag0 = magm[0:15]
|
|
val = var[magm]
|
|
ele = self.om.getElement(mag0)
|
|
if 'mqua' in magm or 'mqsk' in magm:
|
|
if not ele is None:
|
|
ele.k1 = val
|
|
print('Updating %s to k1: %8.4f' % (mag0, val))
|
|
magnets=self.parent.reference.getDependence(mag0)
|
|
if not magnets is None:
|
|
for magd in magnets:
|
|
ele = self.om.getElement(magd)
|
|
if not ele is None:
|
|
ele.k1 = val
|
|
print('Updating %s to k1: %8.4f' % (magd, val))
|
|
elif 'mkac' in magm or 'mkdc' in magm:
|
|
if not ele is None:
|
|
ele.design_kick = val
|
|
ele.cory = val
|
|
print('Updating %s to design_kick: %8.4f' % (mag0, val))
|
|
##################
|
|
# tracking
|
|
|
|
def track(self):
|
|
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]
|
|
refloc, twiss0 = self.parent.reference.getReference()
|
|
|
|
if refloc.upper() == 'START':
|
|
refloc = start.upper()
|
|
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
|
|
print('Tracking from',start,'to',end)
|
|
self.doTrack(start,end,refloc,twiss0)
|
|
|
|
|
|
def doTrack(self, start,end, refloc, twiss_in, plot = True):
|
|
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)
|
|
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('swissfel',start+'$START',end+'$END',twiss0)
|
|
energy = self.calcEnergyProfile(twiss)
|
|
if plot:
|
|
self.parent.plot.newData(twiss,energy)
|
|
|
|
|
|
|
|
|
|
def doBackTrack(self,start=None,end=None,twiss0=None):
|
|
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)
|
|
# print(start,end)
|
|
twiss = self.madx.track('invswissfel', start, end+'$START',twiss0)
|
|
twiss1 = {}
|
|
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 calcEnergyProfile(self,twiss):
|
|
energy = np.array([0. for i in range(len(twiss.betx))])
|
|
e0 = 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
|
|
return energy
|
|
|
|
def setBranch(self,end):
|
|
destination = 'ARAMIS'
|
|
if 'SPO' in end:
|
|
destination = 'PORTHOS'
|
|
elif 'SAT' in end:
|
|
destination = 'ATHOS'
|
|
elif 'S10BD' in end or 'SIN' in end:
|
|
destination = 'INJECTOR'
|
|
self.om.setBranch(destination,'SINLH01')
|
|
self.order=self.om.getBranchElements()
|
|
self.madx.updateLattice(self.om,destination,self.forceLat)
|
|
self.forceLat=False
|
|
|
|
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
|
|
|
|
|