358 lines
14 KiB
Python
358 lines
14 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.ElegantPlotGui import Ui_ElegantPlotGUI
|
|
import numpy as np
|
|
import h5py
|
|
|
|
|
|
|
|
class ElegantPlot(QtWidgets.QMainWindow, Ui_ElegantPlotGUI):
|
|
def __init__(self,parent=None):
|
|
super(ElegantPlot, 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)
|
|
|
|
|
|
|
|
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)
|
|
|
|
|
|
def LPS(self,t,p,Q):
|
|
title = self.parent.UIDistList.currentText()
|
|
x = np.array(t).ravel()*1e12 # to ps
|
|
y = np.array(p).ravel()*0.511 # to MEV
|
|
|
|
|
|
xmin = x.min()
|
|
xmax = x.max()
|
|
ymin = y.min()
|
|
ymax = y.max()
|
|
|
|
|
|
N = 300
|
|
img,xed,yed = np.histogram2d(x,y,N)
|
|
xdist, xval = np.histogram(x,N)
|
|
cal = np.sum(xdist)*(xval[1]-xval[0])*1e-12/Q
|
|
xdist = xdist/cal
|
|
|
|
ydist, yval = np.histogram(y, N)
|
|
ydist = ydist / np.max(ydist) * 0.3 * (xmax - xmin) + xmin
|
|
|
|
self.axes.clear()
|
|
self.axes2.clear()
|
|
self.axes.imshow(np.transpose(np.fliplr(img)), aspect='auto', interpolation='bicubic', cmap='viridis',
|
|
extent=[xmin, xmax, ymin, ymax])
|
|
self.axes2.plot(xval[:-1],xdist,'y')
|
|
self.axes.plot(ydist,yval[:-1], 'c', label='Energy Distribution')
|
|
self.axes.plot([], [], 'y', label='Current')
|
|
self.axes.legend()
|
|
|
|
self.axes.set_xlabel('t (ps)')
|
|
self.axes.set_ylabel('E (MeV)')
|
|
self.axes2.set_ylabel('I (A)')
|
|
self.axes2.set_ylim([0,3*np.max(xdist)])
|
|
self.axes.set_title(title)
|
|
self.canvas.draw()
|
|
|
|
|
|
|
|
# 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.PMux.toggled.connect(self.doplot)
|
|
# self.PMuy.toggled.connect(self.doplot)
|
|
# self.PX.toggled.connect(self.doplot)
|
|
# self.PY.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 clearReference(self):
|
|
self.twissref=None
|
|
self.energyref=None
|
|
self.doplot()
|
|
|
|
def newReference(self):
|
|
self.twissref=self.twiss
|
|
self.energyref=copy.deepcopy(self.energy)
|
|
|
|
def newData(self,twiss,energy):
|
|
if not self.isVisible():
|
|
self.show()
|
|
print('Updating Plotting Data')
|
|
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()
|
|
filt['X'] = self.PX.isChecked()
|
|
filt['Y'] = self.PY.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_{x}$')
|
|
ylabel = ylabel + r'$\mu_x$, '
|
|
if filt['MUY']:
|
|
self.plotSingle(s[i1:i2], self.twiss.muy[i1:i2], (1, 0.5, 0, 1), r'$\mu_{y}$')
|
|
ylabel = ylabel + r'$\mu_y$, '
|
|
if filt['X']:
|
|
self.plotSingle(s[i1:i2], self.twiss.x[i1:i2], (0, 0.5, 1, 1), r'$x$')
|
|
ylabel = ylabel + r'$x$ (m), '
|
|
if filt['Y']:
|
|
self.plotSingle(s[i1:i2], self.twiss.y[i1:i2], (1, 0.5, 0, 1), r'$y$')
|
|
ylabel = ylabel + r'$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}$ (ref)',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}$ (ref)',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}$ (ref)',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}$ (ref)',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}$ (ref)',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}$ (ref)',dashed=True)
|
|
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_{x}$ (ref)',dashed=True)
|
|
ylabel = ylabel + r'$\mu_x$, '
|
|
if filt['MUY']:
|
|
self.plotSingle(s[i1:i2], self.twiss.muy[i1:i2], (1, 0.5, 0, 1), r'$\mu_{y}$ (ref)',dashed=True)
|
|
ylabel = ylabel + r'$\mu_y$, '
|
|
if filt['X']:
|
|
self.plotSingle(s[i1:i2], self.twiss.x[i1:i2], (0, 0.5, 1, 1), r'$x$ (ref)',dashed=True)
|
|
ylabel = ylabel + r'$x$ (m), '
|
|
if filt['Y']:
|
|
self.plotSingle(s[i1:i2], self.twiss.y[i1:i2], (1, 0.5, 0, 1), r'$y$ (ref)',dashed=True)
|
|
ylabel = ylabel + r'$y$ (m), '
|
|
if filt['RE56']:
|
|
self.plotSingle(s[i1:i2], self.twissref.re56[i1:i2], (0, 0, 0, 1), r'$R_{56}$ (ref)',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$ (ref)',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)
|
|
|
|
|
|
|
|
|
|
# --------------------------------
|
|
# 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 = ElegantPlot()
|
|
plot.show()
|
|
sys.exit(app.exec_())
|