317 lines
12 KiB
Python
317 lines
12 KiB
Python
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['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 = 10
|
|
self.UITwissValues.setColumnCount(ncol)
|
|
self.UITwissValues.setRowCount(nrow)
|
|
self.UITwissValues.setHorizontalHeaderLabels(['Name','s','betax','betay','alphax','alphay','etax','etay','r56','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.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()
|
|
|
|
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], 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_())
|