import json import numpy as np from onlinemodel.core import Facility 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(init=1, alt = phase) self.order = None self.madx = CMadX() # hook up events self.eventHandling() def updateEnergy(self,E0): self.om.forceEnergyAt('SINLH02.MBND100', E0*1e6) def updateElement(self,name,val): if 'MQUA' in name: self.om.setRegExpElement(name[0:7], name[8:15], 'k1', float(val[0])) if 'MSEX' in name: self.om.setRegExpElement(name[0:7], name[8:15], 'k2', float(val[0])) 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 self.om.setRegExpElement(name[0:7], 'RACC', 'Gradient', grad) 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={} loc = 'SINLH02.MBND100' energy={'location': loc, 'energy':self.om.EnergyAt(loc)[0]} 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 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} 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]) ##### very old code should be become obsolete after some nice matching files 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]) ##############################33 # tracking def eventHandling(self): self.parent.UITrack.clicked.connect(self.track) 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] 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 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: twiss0 = self.doBackTrack(refloc,start,twiss0) self.madx.updateVariables(twiss0) twiss = self.madx.track(start+'$START',end+'$END') energy = np.array([0. for i in range(len(twiss.betx))]) 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 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' 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) 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