ElegantTools/model.py

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/ReferenceSetting.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=[]):
        start,end=self.checkRange(start,end)
        self.setBranch(end)
        lat=Lattice()
        lat.charge = Q
        lat.dump = dump
        print(dump)
        if lsc:
            lat.lsc=1
        else:
            lat.lsc=0
        if csr:
            lat.csr=1
        else:
            lat.csr=0
        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)
            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'])
        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']
        self.parent.UIDistBetax.setText('%7.3f' % self.startTwiss['betax0'])
        self.parent.UIDistBetay.setText('%7.3f' % self.startTwiss['betay0'])
        self.parent.UIDistAlphax.setText('%7.3f' % self.startTwiss['alphax0'])
        self.parent.UIDistAlphay.setText('%7.3f' % self.startTwiss['alphay0'])
        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))
            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()))*np.pi/180.
            gain=grad*np.sin(phase)
            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
            print(sec,volt,phase)
            self.om.setRegExpElement(sec, 'RACC', 'Gradient', volt)
            self.om.setRegExpElement(sec, 'RACC', 'Phase', phase)

        self.om.forceEnergyAt(self.initEnergyRef, self.initEnergy * 1e6)
        self.parent.UIRFSettings.resizeColumnsToContents()
        self.parent.UIRFSettings.horizontalHeader().show()