Files

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)