Initial Round of Sandbox implementation
This commit is contained in:
316
plot.py
Normal file
316
plot.py
Normal file
@@ -0,0 +1,316 @@
|
||||
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_())
|
||||
Reference in New Issue
Block a user