import sys 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 class OpticsPlot(QtWidgets.QMainWindow, Ui_OpticsPlotGUI): def __init__(self,phase=0): super(OpticsPlot, self).__init__() self.setupUi(self) 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 # 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.PStart.editingFinished.connect(self.doplot) self.PEnd.editingFinished.connect(self.doplot) def newData(self,twiss): self.twiss=twiss self.updateOpticsTable() self.doplot() def updateOpticsTable(self): self.UITwissValues.clear() if self.twiss is None: return nrow = len(self.twiss.s) ncol = 9 self.UITwissValues.setColumnCount(ncol) self.UITwissValues.setRowCount(nrow) self.UITwissValues.setHorizontalHeaderLabels(['Name','s','betax','betay','alphax','alphay','etax','etay','r56']) 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.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() 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['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], data['Energy'][i1:i2], (0, 1, 0, 1), r'$E$') # ylabel = ylabel + r'$E$ (MeV), ' if len(ylabel) < 3: self.canvas.draw() return 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_())