233 lines
8.3 KiB
Python
233 lines
8.3 KiB
Python
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_())
|