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)