198 lines
7.8 KiB
Python
198 lines
7.8 KiB
Python
from threading import Thread
|
|
import subprocess
|
|
import re
|
|
import os
|
|
import sdds
|
|
|
|
from PyQt5 import QtCore
|
|
|
|
class Track(QtCore.QObject):
|
|
|
|
message = QtCore.pyqtSignal(str) # any print out of this class
|
|
done = QtCore.pyqtSignal()
|
|
|
|
def __init__(self, parent=None):
|
|
super(Track, self).__init__()
|
|
|
|
self.parent=parent
|
|
self.isRunning=False
|
|
self.dumpPoints=[]
|
|
self.parent.UITrack.clicked.connect(self.startTracking)
|
|
self.parent.UITrackReload.clicked.connect(self.reloadTracking)
|
|
self.parent.UIAddDiagnosticDump.clicked.connect(self.addDiagnosticDump)
|
|
self.parent.UIAddSectionDump.clicked.connect(self.addSectionDump)
|
|
self.parent.UIDeleteDump.clicked.connect(self.deleteDump)
|
|
self.message.connect(self.trackMessage)
|
|
self.done.connect(self.elegantFinished)
|
|
|
|
def startTracking(self):
|
|
if self.isRunning:
|
|
return
|
|
project = str(self.parent.UITrackRootName.text())
|
|
self.root= 'Runs/'+project
|
|
|
|
start=str(self.parent.UITrackStart.text())
|
|
end=str(self.parent.UITrackEnd.text())
|
|
lsc = self.parent.UITrackDoLSC.isChecked()
|
|
csr = self.parent.UITrackDoCSR.isChecked()
|
|
center = self.parent.UITrackAutoAlign.isChecked()
|
|
if center:
|
|
aligns = ['SINLH03.DBPM010','SINDI01.DWSC090','S10MA01.DBPM010','SATDI01.DBCM010','SATUN04.DBPM010','SARMA01.DBPM040']
|
|
else:
|
|
aligns= [ ]
|
|
print('Tracking from',start,'to',end,'...')
|
|
if self.parent.preferredDist is None:
|
|
Q = 200e-12
|
|
else:
|
|
Q = self.parent.preferredDist.getQ()
|
|
dumps = []
|
|
for ele in self.dumpPoints:
|
|
if '-' in ele:
|
|
dumps.append(ele.replace('-','.'))
|
|
else:
|
|
dumps.append(ele+'.End')
|
|
print(dumps)
|
|
latfile = self.parent.model.writeLattice(project,start,end,lsc,csr,Q,dumps,aligns)
|
|
print('Lattice written to',latfile)
|
|
if self.parent.preferredDist is None:
|
|
print('No distribution defined')
|
|
print('Only lattice file written to Directory ./Runs/')
|
|
return
|
|
distfile='Runs/'+project+'Inputdist.sdds'
|
|
self.parent.preferredDist.saveDist(distfile)
|
|
print('Input distribution copied to',distfile)
|
|
res=self.parent.preferredDist.analyseBeamShort(distfile)
|
|
inputfile=self.writeInputFile('Runs',project,latfile,distfile,res)
|
|
print('Input file written to',inputfile)
|
|
# self.cmd=['mpirun','-np','10','Pelegant',inputfile]
|
|
self.cmd=['elegant',inputfile]
|
|
self.isRunning=True
|
|
self.parent.UITrackStatus.setText('Elegant running...')
|
|
self.parent.UITrackStatus.setStyleSheet("background-color: yellow")
|
|
Thread(name='Elegant Run', target=self.Telegant).start()
|
|
|
|
def Telegant(self):
|
|
self.parent.UITrackOutput.appendPlainText('Starting thread to run elegant')
|
|
# p = subprocess.run(self.cmd)
|
|
p = subprocess.Popen(self.cmd,stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
|
|
for line in iter(p.stdout.readline, b''):
|
|
self.message.emit(line.decode('utf-8').strip())
|
|
self.isRunning=False
|
|
self.done.emit()
|
|
|
|
def reloadTracking(self):
|
|
if self.isRunning:
|
|
return
|
|
project = str(self.parent.UITrackRootName.text())
|
|
self.root = 'Runs/' + project
|
|
self.loadTracking()
|
|
|
|
@QtCore.pyqtSlot()
|
|
def elegantFinished(self):
|
|
self.parent.UITrackStatus.setText('Idle')
|
|
self.parent.UITrackStatus.setStyleSheet("background-color: white")
|
|
self.loadTracking()
|
|
|
|
def loadTracking(self):
|
|
data = sdds.load(self.root+'.cen')
|
|
columns = data.columnName
|
|
res = {}
|
|
fields = ['s','ElementName','Cx','Cxp','Cy','Cyp','Cdelta','pCentral']
|
|
for field in fields:
|
|
res[field] = data.columnData[columns.index(field)][0]
|
|
del data
|
|
data = sdds.load(self.root + '.sig')
|
|
columns = data.columnName
|
|
fields = ['Sx', 'Sxp', 'Sy', 'Syp', 'Sdelta', 'St','enx','eny','ecnx','ecny','betaxBeam','betayBeam','alphaxBeam','alphayBeam']
|
|
for field in fields:
|
|
res[field] = data.columnData[columns.index(field)][0]
|
|
del data
|
|
dist = [self.root+'Inputdist.sdds']+self.getDumps()
|
|
self.parent.plot.newData(res,dist)
|
|
|
|
def getDumps(self):
|
|
filt=re.compile('%s.*out' % self.root.split('/')[-1])
|
|
dumps=[]
|
|
directory = os.fsencode('Runs')
|
|
for file in os.listdir(directory):
|
|
filename = os.fsdecode(file)
|
|
if re.match(filt,filename):
|
|
print('found dump:',filename)
|
|
dumps.append('Runs/'+filename)
|
|
return dumps
|
|
|
|
@QtCore.pyqtSlot(str)
|
|
def trackMessage(self, msg=''):
|
|
# print(msg)
|
|
self.parent.UITrackOutput.appendPlainText(msg)
|
|
|
|
def writeInputFile(self,path,name,latfile,distfile,distinfo):
|
|
filename = '%s/%s.ele' % (path, name)
|
|
with open(filename, 'w') as fid:
|
|
fid.write('&run_setup\n')
|
|
fid.write('\tlattice\t\t= %s,\n' % (latfile))
|
|
fid.write('\tuse_beamline\t= SwissFEL,\n')
|
|
fid.write('\trootname\t= %s,\n' % name)
|
|
fid.write('\toutput\t\t= %s/%%s.out,\n' % path)
|
|
fid.write('\tcentroid\t\t= %s/%%s.cen,\n' % path)
|
|
fid.write('\tsigma\t\t= %s/%%s.sig,\n' % path)
|
|
fid.write('\tfinal\t\t= %s/%%s.fin,\n' % path)
|
|
fid.write('\tparameters\t= %s/%%s.par,\n' % path)
|
|
fid.write('\tmagnets\t\t= %s/%%s.mag,\n' % path)
|
|
fid.write('\tcombine_bunch_statistics = 0,\n')
|
|
fid.write('\tdefault_order\t= 2,\n')
|
|
fid.write('\tconcat_order\t= 0,\n')
|
|
fid.write('\tprint_statistics\t= 0,\n')
|
|
fid.write('\trandom_number_seed\t= 9876543210,\n')
|
|
fid.write('\tp_central\t= %f,\n' % distinfo['pAverage'])
|
|
fid.write('\ttracking_updates\t= 1\n')
|
|
fid.write('\talways_change_p0\t= 1\n')
|
|
fid.write('&end\n\n')
|
|
fid.write('&run_control\n')
|
|
fid.write('\tn_steps\t= 1,\n')
|
|
fid.write('\treset_rf_for_each_step = 1\n')
|
|
fid.write('&end\n\n')
|
|
|
|
fid.write('&twiss_output\n')
|
|
fid.write('\tfilename\t= %s/%%s.twi,\n' % path)
|
|
fid.write('\tmatched\t\t= 0,\n')
|
|
fid.write('\tbeta_x\t\t= %f,\n' % distinfo['betax'])
|
|
fid.write('\tbeta_y\t\t= %f,\n' % distinfo['betay'])
|
|
fid.write('\talpha_x\t\t= %f,\n' % distinfo['alphax'])
|
|
fid.write('\talpha_y\t\t= %f,\n' % distinfo['alphay'])
|
|
fid.write('&end\n\n')
|
|
|
|
fid.write('&sdds_beam\n')
|
|
fid.write('\tinput_type= "elegant",\n')
|
|
fid.write('\tsample_interval\t= 1,\n')
|
|
fid.write('\tinput = %s,\n' % distfile)
|
|
fid.write('\treuse_bunch\t= 0 \n')
|
|
fid.write('&end\n\n')
|
|
|
|
fid.write('&track\n')
|
|
fid.write('&end\n\n')
|
|
fid.close()
|
|
|
|
return filename
|
|
|
|
def addDiagnosticDump(self):
|
|
self.addDump(self.parent.UIDiagnosticList)
|
|
|
|
def addSectionDump(self):
|
|
self.addDump(self.parent.UISectionList)
|
|
|
|
def addDump(self,qwidget):
|
|
elements=[str(ele.text()) for ele in qwidget.selectedItems()]
|
|
for ele in elements:
|
|
if not ele in self.dumpPoints:
|
|
self.dumpPoints.append(ele)
|
|
self.parent.UIDumpLocation.clear()
|
|
for ele in self.dumpPoints:
|
|
self.parent.UIDumpLocation.addItem(ele)
|
|
|
|
def deleteDump(self):
|
|
sels = [str(ele.text()) for ele in self.parent.UIDumpLocation.selectedItems()]
|
|
self.dumpPoints = [ele for ele in self.dumpPoints if not ele in sels]
|
|
self.parent.UIDumpLocation.clear()
|
|
for ele in self.dumpPoints:
|
|
self.parent.UIDumpLocation.addItem(ele) |