SwissFEL-Simulation/ElegantTools
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
main
ElegantTools/plot.py

385 lines
15 KiB
Python
Raw Permalink Blame History

import sys
import pysdds
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.data = None
self.setWindowIcon(QtGui.QIcon("rsc/Audrey_Icon.png"))
self.setWindowTitle("Elegant Results Plotting Window")
self.PBetax.toggled.connect(self.doplot)
self.PBetay.toggled.connect(self.doplot)
self.PSizex.toggled.connect(self.doplot)
self.PSizey.toggled.connect(self.doplot)
self.PEmitx.toggled.connect(self.doplot)
self.PEmity.toggled.connect(self.doplot)
self.PDuration.toggled.connect(self.doplot)
self.PEnergy.toggled.connect(self.doplot)
self.PSpread.toggled.connect(self.doplot)
self.PCEnergy.toggled.connect(self.doplot)
self.PCx.toggled.connect(self.doplot)
self.PCy.toggled.connect(self.doplot)
self.PDistTE.clicked.connect(self.doPSplot)
self.PDistXPx.clicked.connect(self.doPSplot)
self.PDistYPy.clicked.connect(self.doPSplot)
self.PDistTX.clicked.connect(self.doPSplot)
self.PDistTY.clicked.connect(self.doPSplot)
self.PDistXY.clicked.connect(self.doPSplot)
self.currentDist=None
self.dist={}
self.Q=0
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 newData(self,data,dists=[]):
if not self.isVisible():
self.show()
print('Updating Plotting Data')
self.data=data
# scale some data
self.data['ecnx']=[ele*1e9 for ele in self.data['ecnx']] # nm
self.data['ecny'] = [ele * 1e9 for ele in self.data['ecny']] # nm
self.data['Sx'] = [ele * 1e6 for ele in self.data['Sx']] # micrometer
self.data['Sy'] = [ele * 1e6 for ele in self.data['Sy']] # micrometer
self.data['Cx'] = [ele * 1e6 for ele in self.data['Cx']] # micrometer
self.data['Cy'] = [ele * 1e6 for ele in self.data['Cy']] # micrometer
self.data['pCentral'] = [ele * 0.511 for ele in self.data['pCentral']] # MeV
self.data['St'] = [ele * 1e15 for ele in self.data['St']] # fs
self.data['Sdelta'] = [ele * 1e2 for ele in self.data['Sdelta']] # %
self.data['Cdelta'] = [ele * 1e2 for ele in self.data['Cdelta']] # %
self.PDistList.clear()
for dist in dists:
self.PDistList.addItem(dist)
self.currentDist=None
self.updateOpticsTable()
self.doplot()
def updateOpticsTable(self):
self.UITwissValues.clear()
if self.data is None:
return
nrow = len(self.data['s'])
ncol = 14
self.UITwissValues.setColumnCount(ncol)
self.UITwissValues.setRowCount(nrow)
self.UITwissValues.setHorizontalHeaderLabels(['Element','s','Size in x','Size in y','Size in t','Energy','Energy Spread','Emit x','Emit y','Betax','Betay','<x>','<y>','<E>'])
for i in range(nrow):
self.UITwissValues.setItem(i, 0, QtWidgets.QTableWidgetItem(self.data['ElementName'][i]))
self.UITwissValues.setItem(i, 1, QtWidgets.QTableWidgetItem('%10.6f' % self.data['s'][i]))
self.UITwissValues.setItem(i, 2, QtWidgets.QTableWidgetItem('%10.6f' % self.data['Sx'][i]))
self.UITwissValues.setItem(i, 3, QtWidgets.QTableWidgetItem('%10.6f' % self.data['Sy'][i]))
self.UITwissValues.setItem(i, 4, QtWidgets.QTableWidgetItem('%10.6f' % self.data['St'][i]))
self.UITwissValues.setItem(i, 5, QtWidgets.QTableWidgetItem('%10.6f' % self.data['pCentral'][i]))
self.UITwissValues.setItem(i, 6, QtWidgets.QTableWidgetItem('%10.6f' % self.data['Sdelta'][i]))
self.UITwissValues.setItem(i, 7, QtWidgets.QTableWidgetItem('%10.6f' % self.data['ecnx'][i]))
self.UITwissValues.setItem(i, 8, QtWidgets.QTableWidgetItem('%10.6f' % self.data['ecny'][i]))
self.UITwissValues.setItem(i, 9, QtWidgets.QTableWidgetItem('%10.6f' % self.data['betaxBeam'][i]))
self.UITwissValues.setItem(i, 10, QtWidgets.QTableWidgetItem('%10.6f' % self.data['betayBeam'][i]))
self.UITwissValues.setItem(i, 11, QtWidgets.QTableWidgetItem('%10.6f' % self.data['Cx'][i]))
self.UITwissValues.setItem(i, 12, QtWidgets.QTableWidgetItem('%10.6f' % self.data['Cy'][i]))
self.UITwissValues.setItem(i, 13, QtWidgets.QTableWidgetItem('%10.6f' % self.data['Cdelta'][i]))
self.UITwissValues.resizeColumnsToContents()
self.UITwissValues.verticalHeader().hide()
self.UITwissValues.setEditTriggers(QtWidgets.QAbstractItemView.NoEditTriggers)
def loadDist(self,filename):
dist = pysdds.read(filename)
cols=['t','p','x','xp','y','yp']
self.dist.clear()
for col in cols:
self.dist[col] = dist.col(col)[0]
self.Q = dist.par('Charge').data[0]
return filename
def doPSplot(self):
dist = self.PDistList.currentText()
if not dist == self.currentDist:
self.currentDist = self.loadDist(dist)
but = self.sender().objectName()
if 'DistTE' in but:
self.LPS(self.dist['t'],self.dist['p'],self.Q,dist)
elif 'DistXPx' in but:
self.TPS(self.dist['x'],self.dist['xp'],True,dist)
elif 'DistYPy' in but:
self.TPS(self.dist['y'], self.dist['yp'], False,dist)
elif 'DistTX' in but:
self.TimePS(self.dist['t'], self.dist['x'], self.Q,True,dist)
elif 'DistTY' in but:
self.TimePS(self.dist['t'], self.dist['y'], self.Q,False,dist)
elif 'DistXY' in but:
self.TS(self.dist['x'], self.dist['y'],dist)
def LPS(self,t,p,Q,title=None):
if title is None:
title = self.parent.UIDistList.currentText()
x = np.array(t).ravel()*1e15 # to fs
y = np.array(p).ravel()*0.511 # to MEV
x=x-np.mean(x)
self.PS(x,y,Q=Q,title = title, labx=r'$t$ (fs)',laby=r'$E$ (MeV)',laby2=r'$I$ (A)',legend1='Energy Distribution',legend2='Current')
def TPS(self,xin,pin,isX=True,title=None):
if title is None:
title = self.parent.UIDistList.currentText()
x = np.array(xin).ravel()*1e3 # to mm
y = np.array(pin).ravel()*1e3 # to mrad
labx=r'$y$ (mm)'
laby=r'$y\prime$ (mrad)'
if isX:
labx = r'$x$ (mm)'
laby = r'$x\prime$ (mrad)'
self.PS(x, y, Q=None, title=title, labx=labx, laby=laby, laby2='',legend1='Distribution in Momentum', legend2='Distribution in Space')
def TimePS(self,t,xin,Q,isX=True,title=None):
if title is None:
title = self.parent.UIDistList.currentText()
x = np.array(t).ravel()*1e15 # to fs
y = np.array(xin).ravel()*1e3 # to mm
x=x-np.mean(x)
labx=r'$t$ (fs)'
laby=r'$y$ (mm)'
if isX:
laby = r'$x$ (m)'
self.PS(x, y, Q=Q, title=title, labx=labx, laby=laby, laby2=r'$I$ (A)',legend1='Distribution in Space', legend2='Current')
def TS(self,xin,yin,title=None):
if title is None:
title = self.parent.UIDistList.currentText()
x = np.array(xin).ravel()*1e3 # to mm
y = np.array(yin).ravel()*1e3 # to mm
labx=r'$x$ (mm)'
laby=r'$y$ (mm)'
self.PS(x, y, Q=None, title=title, labx=labx, laby=laby, laby2='',legend1='Distribution in Y', legend2='Distribution in X')
def PS(self,x,y,Q = None,title='',labx='',laby='',laby2='',legend1='',legend2=''):
xmin = x.min()
xmax = x.max()
ymin = y.min()
ymax = y.max()
dx=(xmax-xmin)*0.1
dy=(ymax-ymin)*0.1
xmax+=dx
xmin-=dx
ymax+=dy
ymin-=dy
range=[[xmin,xmax],[ymin,ymax]]
N = 300
img,xed,yed = np.histogram2d(x,y,N,range=range)
xdist, xval = np.histogram(x,N,range=range[0])
if Q:
cal = np.sum(xdist)*(xval[1]-xval[0])*1e-12/Q
xdist = xdist/cal*1e3
ydist, yval = np.histogram(y, N,range=range[1])
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', cmap='gnuplot2',interpolation='bicubic', # was viridis andf
extent=[xmin, xmax, ymin, ymax])
self.axes2.plot(xval[:-1],xdist,'y')
self.axes.plot(ydist,yval[:-1], 'c', label=legend1)
self.axes.plot([], [], 'y', label=legend2)
self.axes.legend()
self.axes.set_xlabel(labx)
self.axes.set_ylabel(laby)
if len(laby2)> 1:
self.axes2.yaxis.set_visible(True)
else:
self.axes2.yaxis.set_visible(False)
self.axes2.set_ylim([0,3*np.max(xdist)])
self.axes2.set_xlim([xmin,xmax])
self.axes2.set_ylabel(laby2)
self.axes.set_title(title)
self.canvas.draw()
def doplot(self):
if self.data 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['Sx']=self.PSizex.isChecked()
filt['Sy'] = self.PSizey.isChecked()
filt['ecnx']=self.PEmitx.isChecked()
filt['ecny']=self.PEmity.isChecked()
filt['betaxBeam']=self.PBetax.isChecked()
filt['betayBeam']=self.PBetay.isChecked()
filt['Cx']=self.PCx.isChecked()
filt['Cy']=self.PCy.isChecked()
filt['St']=self.PDuration.isChecked()
filt['pCentral']=self.PEnergy.isChecked()
filt['Sdelta'] = self.PSpread.isChecked()
filt['Cdelta'] = self.PCEnergy.isChecked()
color={}
color['Sx']={'color':(0,0,1,1),'label':r'$\sigma_x$ ($\mu$m)','legend':r'$\sigma_x$','dash':False}
color['Sy'] = {'color':(1, 0, 0, 1),'label':r'$\sigma_y$ ($\mu$m)','legend':r'$\sigma_y$','dash':False}
color['ecnx'] ={'color': (0, 0.4, 1, 1),'label':r'$\epsilon_{nx}$ (nm)','legend':r'$\epsilon_{nx}$','dash':False}
color['ecny'] = {'color':(1, 0.4, 0, 1),'label':r'$\epsilon_{ny}$ (nm)','legend':r'$\epsilon_{ny}$','dash':False}
color['betaxBeam'] = {'color': (0, 0, 1, 1), 'label': r'$\beta_x$ (m)', 'legend': r'$\beta_x$','dash':False}
color['betayBeam'] = {'color': (1, 0, 0, 1), 'label': r'$\beta_y$ (m)', 'legend': r'$\beta_y$','dash':False}
color['Cx'] = {'color': (0, 0, 1, 1), 'label': r'$<x>$ ($\mu$m)', 'legend': r'$<x>$','dash':True}
color['Cy'] = {'color': (1, 0, 0, 1), 'label': r'$<y>$ ($\mu$m)', 'legend': r'$<y>$','dash':True}
color['pCentral'] = {'color':(0,0.5,0,1),'label':r'$E_0$ (MeV)', 'legend': r'$E_0$','dash':False}
color['Sdelta'] = {'color': (0, 0.5, 0, 1), 'label': r'$\sigma_\delta$ (%)', 'legend': r'$\sigma_\delta$', 'dash': False}
color['Cdelta'] = {'color': (0, 0.5, 0, 1), 'label': r'$<\delta>$ (%)', 'legend': r'$<\delta>$', 'dash': True}
color['St'] = {'color': (0.5, 0, 0.5, 1), 'label': r'$\sigma_t$ (fs)', 'legend': r'$\sigma_t$', 'dash': False}
s = np.array(self.data['s'])
i1 = np.argmin(np.abs(s - z0))
i2 = np.argmin(np.abs(s - z1))
self.axes.clear()
self.axes2.clear()
ylabel = r''
for key in filt.keys():
if filt[key]:
if key in color.keys():
self.plotSingle(s[i1:i2],self.data[key][i1:i2], color[key]['color'],color[key]['legend'],dashed=color[key]['dash'])
ylabel += color[key]['label']+', '
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.data['ElementName'])
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 and '.B' 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 and '.SEX':
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 and '.WIG' 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 '.RACC' in name or '.RTDS' 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 and '.Q':
if splitquads == True:
if '.Q2' 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 '.Q1' 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_())
Reference in New Issue View Git Blame Copy Permalink