import sys import copy from PyQt5 import QtWidgets,QtGui from matplotlib.figure import Figure import matplotlib.patches as patches from matplotlib.backends.backend_qt5agg import ( FigureCanvasQTAgg as FigureCanvas, NavigationToolbar2QT as NavigationToolbar) from ui.OpticsPlotGui import Ui_OpticsPlotGUI import numpy as np import h5py class OpticsPlot(QtWidgets.QMainWindow, Ui_OpticsPlotGUI): def __init__(self,parent=None): super(OpticsPlot, self).__init__() self.setupUi(self) self.parent=parent self.version = '1.0.1' self.setWindowIcon(QtGui.QIcon("rsc/iconoptics.png")) self.setWindowTitle("SwissFEL Optics Plotting Window") self.initmpl(self.mplvl, self.mplwindow) self.twiss=None self.twissref = None self.energy=None # all action for optics plotting self.PBetax.toggled.connect(self.doplot) self.PAlphax.toggled.connect(self.doplot) self.PBetay.toggled.connect(self.doplot) self.PAlphay.toggled.connect(self.doplot) self.PEtax.toggled.connect(self.doplot) self.PEtay.toggled.connect(self.doplot) self.PR56.toggled.connect(self.doplot) self.PEnergy.toggled.connect(self.doplot) self.PStart.editingFinished.connect(self.doplot) self.PEnd.editingFinished.connect(self.doplot) self.UIPlotExportOptics.clicked.connect(self.exportOptics) self.UIPLotNewReference.clicked.connect(self.saveReference) self.UIPlotSaveReference.clicked.connect(self.newReference) self.UIPlotClearReference.clicked.connect(self.clearReference) def saveReference(self): irow = self.UITwissValues.currentRow() if irow < 0: return twiss={} col = ['betax', 'betay', 'alphax', 'alphay', 'etax', 'etay'] for i, col in enumerate(col): twiss[col] = float(str(self.UITwissValues.item(irow,i+2).text())) location = str(self.UITwissValues.item(irow,0).text()) self.parent.reference.swithToUserDefinedLocation() self.parent.reference.updateReferenceLocation(location) self.parent.reference.updateReferencePoint(twiss) def exportOptics(self): options = QtWidgets.QFileDialog.Options() options |= QtWidgets.QFileDialog.DontUseNativeDialog fileName, _ = QtWidgets.QFileDialog.getSaveFileName(self, "Save Optics","optics.h5","HDF5 Files (*.h5)", options=options) if not fileName: return data={} data['s']=np.array(self.twiss.s) data['betax']=np.array(self.twiss.betx) data['betay'] = np.array(self.twiss.bety) data['alphax'] = np.array(self.twiss.alfx) data['alphay'] = np.array(self.twiss.alfy) data['etax'] = np.array(self.twiss.dx) data['etay'] = np.array(self.twiss.dy) data['r56'] = np.array(self.twiss.re56) data['mux'] = np.array(self.twiss.mux) data['muy'] = np.array(self.twiss.muy) data['x'] = np.array(self.twiss.x) data['y'] = np.array(self.twiss.y) data['energy'] = np.array(self.energy) # data['names'] = np.array(self.twiss.name with h5py.File(fileName, 'w') as f: for keys in data.keys(): f.create_dataset(keys, data=data[keys]) def clearReference(self): self.twissref=None self.energyref=None self.doplot() def newReference(self): self.twissref=copy.deepcopy(self.twiss) self.energyref=copy.deepcopy(self.energy) def newData(self,twiss,energy): self.twiss=twiss self.energy=energy self.updateOpticsTable() self.doplot() def updateOpticsTable(self): self.UITwissValues.clear() if self.twiss is None: return nrow = len(self.twiss.s) ncol = 12 self.UITwissValues.setColumnCount(ncol) self.UITwissValues.setRowCount(nrow) self.UITwissValues.setHorizontalHeaderLabels(['Name','s','betax','betay','alphax','alphay','etax','etay','r56','mux','muy','energy']) for i in range(nrow): self.UITwissValues.setItem(i, 0, QtWidgets.QTableWidgetItem(self.twiss.name[i].split(':')[0])) self.UITwissValues.setItem(i, 1, QtWidgets.QTableWidgetItem('%10.6f' % self.twiss.s[i])) self.UITwissValues.setItem(i, 2, QtWidgets.QTableWidgetItem('%10.6f' % self.twiss.betx[i])) self.UITwissValues.setItem(i, 3, QtWidgets.QTableWidgetItem('%10.6f' % self.twiss.bety[i])) self.UITwissValues.setItem(i, 4, QtWidgets.QTableWidgetItem('%10.6f' % self.twiss.alfx[i])) self.UITwissValues.setItem(i, 5, QtWidgets.QTableWidgetItem('%10.6f' % self.twiss.alfy[i])) self.UITwissValues.setItem(i, 6, QtWidgets.QTableWidgetItem('%10.6f' % self.twiss.dx[i])) self.UITwissValues.setItem(i, 7, QtWidgets.QTableWidgetItem('%10.6f' % self.twiss.dy[i])) self.UITwissValues.setItem(i, 8, QtWidgets.QTableWidgetItem('%10.6f' % self.twiss.re56[i])) self.UITwissValues.setItem(i, 9, QtWidgets.QTableWidgetItem('%10.6f' % self.twiss.mux[i])) self.UITwissValues.setItem(i,10, QtWidgets.QTableWidgetItem('%10.6f' % self.twiss.muy[i])) self.UITwissValues.setItem(i,11, QtWidgets.QTableWidgetItem('%10.6f' % self.energy[i])) self.UITwissValues.resizeColumnsToContents() self.UITwissValues.verticalHeader().hide() self.UITwissValues.setEditTriggers(QtWidgets.QAbstractItemView.NoEditTriggers) def doplot(self): if self.twiss is None: return z0=float(str(self.PStart.text())) z1=float(str(self.PEnd.text())) if z0 > z1: tmp = z1 z1 = z0 z0 = tmp filt={} filt['BETX']=self.PBetax.isChecked() filt['BETY']=self.PBetay.isChecked() filt['ALFX']=self.PAlphax.isChecked() filt['ALFY']=self.PAlphay.isChecked() filt['DX']=self.PEtax.isChecked() filt['DY']=self.PEtay.isChecked() filt['RE56']=self.PR56.isChecked() filt['Energy']=self.PEnergy.isChecked() filt['MUX'] = self.PMux.isChecked() filt['MUY'] = self.PMuy.isChecked() s = self.twiss.s i1 = np.argmin(np.abs(s - z0)) i2 = np.argmin(np.abs(s - z1)) self.axes.clear() self.axes2.clear() ylabel = r'' if filt['BETX']: self.plotSingle(s[i1:i2], self.twiss.betx[i1:i2], (0, 0, 1, 1), r'$\beta_{x}$') ylabel = ylabel + r'$\beta_x$ (m), ' if filt['BETY']: self.plotSingle(s[i1:i2], self.twiss.bety[i1:i2], (1, 0, 0, 1), r'$\beta_{y}$') ylabel = ylabel + r'$\beta_y$ (m), ' if filt['ALFX']: self.plotSingle(s[i1:i2], self.twiss.alfx[i1:i2], (0, 0, 1, 1), r'$\alpha_{x}$') ylabel = ylabel + r'$\alpha_x$ (rad), ' if filt['ALFY']: self.plotSingle(s[i1:i2], self.twiss.alfy[i1:i2], (1, 0, 0, 1), r'$\alpha_{y}$') ylabel = ylabel + r'$\alpha_y$ (rad), ' if filt['DX']: self.plotSingle(s[i1:i2], self.twiss.dx[i1:i2], (0, 0, 1, 1), r'$\eta_{x}$') ylabel = ylabel + r'$\eta_x$ (m), ' if filt['DY']: self.plotSingle(s[i1:i2], self.twiss.dy[i1:i2], (1, 0, 0, 1), r'$\eta_{y}$') ylabel = ylabel + r'$\eta_y$ (m), ' if filt['MUX']: self.plotSingle(s[i1:i2], self.twiss.mux[i1:i2], (0, 0.5, 1, 1), r'$\mu_{y}$') if filt['MUY']: self.plotSingle(s[i1:i2], self.twiss.muy[i1:i2], (1, 0.5, 0, 1), r'$\mu_{y}$') if filt['RE56']: self.plotSingle(s[i1:i2], self.twiss.re56[i1:i2], (0, 0, 0, 1), r'$R_{56}$') ylabel = ylabel + r'$R_{56}$ (m), ' if filt['Energy']: self.plotSingle(s[i1:i2], self.energy[i1:i2], (0, 1, 0, 1), r'$E$') ylabel = ylabel + r'$E$ (MeV), ' if len(ylabel) < 3: self.canvas.draw() return if not self.twissref is None: if filt['BETX']: self.plotSingle(s[i1:i2], self.twissref.betx[i1:i2], (0, 0, 1, 1), r'$\beta_{x}$',dashed=True) ylabel = ylabel + r'$\beta_x$ (m), ' if filt['BETY']: self.plotSingle(s[i1:i2], self.twissref.bety[i1:i2], (1, 0, 0, 1), r'$\beta_{y}$',dashed=True) ylabel = ylabel + r'$\beta_y$ (m), ' if filt['ALFX']: self.plotSingle(s[i1:i2], self.twissref.alfx[i1:i2], (0, 0, 1, 1), r'$\alpha_{x}$',dashed=True) ylabel = ylabel + r'$\alpha_x$ (rad), ' if filt['ALFY']: self.plotSingle(s[i1:i2], self.twissref.alfy[i1:i2], (1, 0, 0, 1), r'$\alpha_{y}$',dashed=True) ylabel = ylabel + r'$\alpha_y$ (rad), ' if filt['DX']: self.plotSingle(s[i1:i2], self.twissref.dx[i1:i2], (0, 0, 1, 1), r'$\eta_{x}$',dashed=True) ylabel = ylabel + r'$\eta_x$ (m), ' if filt['DY']: self.plotSingle(s[i1:i2], self.twissref.dy[i1:i2], (1, 0, 0, 1), r'$\eta_{y}$',dashed=True) ylabel = ylabel + r'$\eta_y$ (m), ' if filt['RE56']: self.plotSingle(s[i1:i2], self.twissref.re56[i1:i2], (0, 0, 0, 1), r'$R_{56}$',dashed=True) ylabel = ylabel + r'$R_{56}$ (m), ' if filt['Energy']: self.plotSingle(s[i1:i2], self.energyref[i1:i2], (0, 1, 0, 1), r'$E$',dashed=True) ylabel = ylabel + r'$E$ (MeV), ' self.axes.legend(bbox_to_anchor=(0.15, 0.85)) self.axes.set_xlabel('s (m)') self.axes.set_ylabel(ylabel[0:-2]) self.plotLayout(s,self.twiss.name) self.axes.set_xlim([s[i1], s[i2]]) ylim = self.axes.get_ylim() dl = np.abs(ylim[1] - ylim[0]) yl = [ylim[0], ylim[1] + 0.2 * dl] self.axes.set_ylim(yl) self.axes2.set_xlim([s[i1], s[i2]]) self.canvas.draw() return def plotLayout(self, s, elements): splitquads = False sstart = 0 s1 = [np.min(s), np.max(s)] s2 = [0.9, 0.9] self.axes2.plot(s1, s2, 'k') for i, name in enumerate(elements): if 'mbnd' in name: s1 = s[i - 1] s2 = s[i] self.axes2.add_patch(patches.Rectangle((s1, 0.9), (s2 - s1), 0.03, facecolor='blue', edgecolor="none")) if 'msex' in name: s1 = s[i - 1] s2 = s[i] self.axes2.add_patch( patches.Rectangle((s1, 0.87), (s2 - s1), 0.06, facecolor='green', edgecolor="none")) if 'uind' in name: s1 = s[i - 1] s2 = s[i] self.axes2.add_patch( patches.Rectangle((s1, 0.88), (s2 - s1), 0.04, facecolor='purple', edgecolor="none")) if 'acc' in name or 'tds' in name: s1 = s[i - 1] s2 = s[i] self.axes2.add_patch(patches.Rectangle((s1, 0.89), (s2 - s1), 0.02, facecolor='cyan', edgecolor="none")) if 'mqua' in name: if splitquads == True: if 'end' in name: s1 = sstart s2 = s[i] self.axes2.add_patch( patches.Rectangle((s1, 0.85), (s2 - s1), 0.1, facecolor='red', edgecolor="none")) splitquads = False else: if 'start' in name: splitquads = True sstart = s[i] else: s1 = s[i - 1] s2 = s[i] self.axes2.add_patch( patches.Rectangle((s1, 0.85), (s2 - s1), 0.1, facecolor='red', edgecolor="none")) self.axes2.set_ylim([0, 1]) self.axes2.yaxis.set_visible(False) return def plotSingle(self, x, y, color, legend, dashed=False): if dashed: self.axes.plot(x, y, '--', color=color, label=legend) else: self.axes.plot(x, y, color=color, label=legend) def initmpl(self,mplvl,mplwindow): self.fig=Figure() self.axes=self.fig.add_subplot(111) self.axes2 = self.axes.twinx() self.canvas = FigureCanvas(self.fig) mplvl.addWidget(self.canvas) self.canvas.draw() self.toolbar=NavigationToolbar(self.canvas,mplwindow, coordinates=True) mplvl.addWidget(self.toolbar) # -------------------------------- # Main routine if __name__ == '__main__': QtWidgets.QApplication.setStyle(QtWidgets.QStyleFactory.create("plastique")) app = QtWidgets.QApplication(sys.argv) if len(sys.argv) > 1: arg=int(sys.argv[1]) else: arg=0 plot = OpticsPlot() plot.show() sys.exit(app.exec_())