197 lines
8.3 KiB
Python
197 lines
8.3 KiB
Python
import json
|
|
import numpy as np
|
|
from PyQt5 import QtCore, QtGui, QtWidgets
|
|
from onlinemodel.core import Facility
|
|
from lattice import Lattice
|
|
|
|
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.defineDumpPoints()
|
|
self.modelPath=''
|
|
self.initModel('OnlineModelLattices/EEHGSetting.json')
|
|
self.parent.UIUpdateRFSettings.clicked.connect(self.updateRF)
|
|
|
|
|
|
def writeLattice(self,project='SwissFEL',start='SINLH01',end='SARBD01',lsc=True,csr=True,Q=200e-12,dump=[],align=[]):
|
|
start,end=self.checkRange(start,end)
|
|
self.setBranch(end)
|
|
lat=Lattice()
|
|
lat.charge = Q
|
|
lat.dump = dump
|
|
if lsc:
|
|
lat.lsc=1
|
|
else:
|
|
lat.lsc=0
|
|
if csr:
|
|
lat.csr=1
|
|
else:
|
|
lat.csr=0
|
|
lat.center=align
|
|
filename = lat.openLatticeStream('Runs', project)
|
|
self.om.exportBranch(lat)
|
|
lat.closeLatticeStream()
|
|
return filename
|
|
|
|
def checkRange(self,start,end):
|
|
if self.om.getSection(start) is None:
|
|
print('Invalid staring point for tracking. Setting to SINLH01')
|
|
start='SINLH01'
|
|
if self.om.getSection(end) is None:
|
|
print('Invalid end point for tracking. Setting to SARBD02')
|
|
end= 'SARBD01'
|
|
return start,end
|
|
|
|
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',end)
|
|
self.order=self.om.getBranchElements()
|
|
|
|
|
|
def initModel(self, modelPath):
|
|
self.modelPath = modelPath
|
|
print('Updating model with settings:',self.modelPath)
|
|
with open(modelPath, 'r', encoding='utf-8') as f:
|
|
settings = json.load(f)
|
|
self.loadSettings(settings)
|
|
|
|
|
|
def loadSettingsGroup(self,group,fields,normalized=False):
|
|
for key in group.keys():
|
|
ele = self.om.getElement(key)
|
|
if ele is None:
|
|
print('Cannot find element',key)
|
|
continue
|
|
for field in fields:
|
|
ele.__dict__[field]=group[key][field]
|
|
if normalized:
|
|
ele.__dict__[field]/=ele.Length
|
|
|
|
def loadSettings(self,settings):
|
|
|
|
if not 'Phase' in settings.keys():
|
|
return False
|
|
if not settings['Phase'] == self.phase:
|
|
return False
|
|
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'])
|
|
E0 = settings['Energy']['energy']
|
|
if isinstance(E0, list):
|
|
E0 = E0[0]
|
|
E0Ref = settings['Energy']['location']
|
|
if isinstance(E0, list):
|
|
E0 = E0[0]
|
|
self.initEnergy = E0
|
|
self.initEnergyRef = E0Ref
|
|
self.om.forceEnergyAt(E0Ref, E0 * 1e6)
|
|
self.startTwiss = settings['InitialCondition']
|
|
print('Initial Condition:',self.startTwiss)
|
|
self.parent.UIDistBetax.setText('%7.3f' % self.startTwiss['betx'])
|
|
self.parent.UIDistBetay.setText('%7.3f' % self.startTwiss['bety'])
|
|
self.parent.UIDistAlphax.setText('%7.3f' % self.startTwiss['alfx'])
|
|
self.parent.UIDistAlphay.setText('%7.3f' % self.startTwiss['alfy'])
|
|
self.parent.UIDistEnergy.setText('%7.3f' % E0)
|
|
self.updateRFTable(settings['RF'])
|
|
|
|
def defineDumpPoints(self):
|
|
self.parent.UISectionList.clear()
|
|
secs=self.om.listSection('*')
|
|
for sec in secs:
|
|
self.parent.UISectionList.addItem(sec)
|
|
|
|
self.parent.UIDiagnosticList.clear()
|
|
digs=self.om.listElement('DBPM')
|
|
for dig in digs:
|
|
self.parent.UIDiagnosticList.addItem(dig.replace('.','-'))
|
|
|
|
def updateRFTable(self,settings):
|
|
secs = ['SINSB0[34]','SINXB','S10CB','S20CB','S30CB','SATCB']
|
|
for j,sec in enumerate(secs):
|
|
grad = self.om.getRegExpElement(sec,'RACC','Gradient')
|
|
L = self.om.getRegExpElement(sec, 'RACC', 'Length')
|
|
Phi = self.om.getRegExpElement(sec, 'RACC', 'Phase')
|
|
vec = 0*1j
|
|
for i,grad0 in enumerate(grad):
|
|
phase = Phi[i]*np.pi/180
|
|
vec += grad0*L[i]*np.exp(1j*phase)*1e-6
|
|
phase=np.mod(np.angle(vec),2*np.pi)*180/np.pi
|
|
self.parent.UIRFSettings.setItem(j,1,QtWidgets.QTableWidgetItem('%7.1f' % np.abs(vec)))
|
|
self.parent.UIRFSettings.setItem(j, 2, QtWidgets.QTableWidgetItem('%7.2f' % phase))
|
|
self.parent.UIRFSettings.setItem(j, 3, QtWidgets.QTableWidgetItem('%7.1f' % vec.imag))
|
|
secs = ['SINDI','S30CB','SATMA']
|
|
for j,sec in enumerate(secs):
|
|
grad = self.om.getRegExpElement(sec,'RTDS','Gradient')
|
|
L = self.om.getRegExpElement(sec, 'RTDS', 'Length')
|
|
Phi = self.om.getRegExpElement(sec, 'RTDS', 'Phase')
|
|
vec = 0
|
|
for i,grad0 in enumerate(grad):
|
|
vec += grad0*L[i]*1e-6
|
|
self.parent.UIRFSettings.setItem(j+6,1,QtWidgets.QTableWidgetItem('%7.1f' % vec))
|
|
self.parent.UIRFSettings.setItem(j+6, 2, QtWidgets.QTableWidgetItem('%7.2f' % Phi[0]))
|
|
self.parent.UIRFSettings.setItem(j+6, 3, QtWidgets.QTableWidgetItem('---'))
|
|
self.parent.UIRFSettings.resizeColumnsToContents()
|
|
self.parent.UIRFSettings.horizontalHeader().show()
|
|
|
|
def updateRF(self):
|
|
secs = ['SINSB0[34]', 'SINXB', 'S10CB', 'S20CB', 'S30CB', 'SATCB']
|
|
for i,sec in enumerate(secs):
|
|
grad=float(str(self.parent.UIRFSettings.item(i,1).text()))
|
|
phase = float(str(self.parent.UIRFSettings.item(i, 2).text()))
|
|
gain=grad*np.sin(phase*np.pi/180.)
|
|
self.parent.UIRFSettings.setItem(i, 3, QtWidgets.QTableWidgetItem('%7.2f' % gain))
|
|
L = self.om.getRegExpElement(sec, 'RACC', 'Length')
|
|
nl=len(L)
|
|
volt=grad/nl/L[0]*1e6
|
|
self.om.setRegExpElement(sec, 'RACC', 'Gradient', volt)
|
|
self.om.setRegExpElement(sec, 'RACC', 'Phase', phase)
|
|
|
|
self.om.forceEnergyAt(self.initEnergyRef, self.initEnergy * 1e6)
|
|
secs = ['SINDI', 'S30CB', 'SATMA']
|
|
for i,sec in enumerate(secs):
|
|
grad=float(str(self.parent.UIRFSettings.item(i+6,1).text()))
|
|
phase = float(str(self.parent.UIRFSettings.item(i+6, 2).text()))
|
|
self.parent.UIRFSettings.setItem(i+6, 3, QtWidgets.QTableWidgetItem('---'))
|
|
L = self.om.getRegExpElement(sec, 'RTDS', 'Length')
|
|
nl=len(L)
|
|
volt=grad/nl/L[0]*1e6
|
|
self.om.setRegExpElement(sec, 'RTDS', 'Gradient', volt)
|
|
self.om.setRegExpElement(sec, 'RTDS', 'Phase', phase)
|
|
|
|
self.parent.UIRFSettings.resizeColumnsToContents()
|
|
self.parent.UIRFSettings.horizontalHeader().show()
|
|
|
|
def saveLatticeSettings(self,filename):
|
|
data = {}
|
|
data['Dumps'] = self.parent.track.dumpPoints
|
|
data['Gradient'] = [str(self.parent.UIRFSettings.item(i,1).text()) for i in range(9)]
|
|
data['Phase'] = [str(self.parent.UIRFSettings.item(i, 2).text()) for i in range(9)]
|
|
|
|
with open(filename, 'w', encoding='utf-8') as f:
|
|
json.dump(data,f,indent=4)
|
|
|
|
def loadLatticeSettings(self, filename):
|
|
with open(filename, 'r') as f:
|
|
data=json.load(f)
|
|
for i in range(9):
|
|
self.parent.UIRFSettings.setItem(i, 1, QtWidgets.QTableWidgetItem(data['Gradient'][i]))
|
|
self.parent.UIRFSettings.setItem(i, 2, QtWidgets.QTableWidgetItem(data['Phase'][i]))
|
|
self.updateRF()
|
|
|
|
self.parent.track.dumpPoints=data['Dumps']
|
|
self.parent.UIDumpLocation.clear()
|
|
for ele in data['Dumps']:
|
|
self.parent.UIDumpLocation.addItem(ele) |