158 lines
4.7 KiB
Python
158 lines
4.7 KiB
Python
#!/opt/gfa/python-3.5/latest/bin/python
|
|
|
|
import sys
|
|
import datetime
|
|
import os.path
|
|
import time
|
|
from os import walk
|
|
import numpy as np
|
|
import sys
|
|
import re
|
|
|
|
|
|
from matplotlib.figure import Figure
|
|
import matplotlib.patches as patches
|
|
|
|
from matplotlib.backends.backend_qt5agg import (
|
|
FigureCanvasQTAgg as FigureCanvas,
|
|
NavigationToolbar2QT as NavigationToolbar)
|
|
|
|
|
|
|
|
|
|
class OpticsPlotting:
|
|
|
|
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 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 plot(self,data,filt,z0,z1):
|
|
|
|
self.axes.clear()
|
|
self.axes2.clear()
|
|
|
|
s=np.array(data['S'])
|
|
if z0>z1:
|
|
tmp=z1
|
|
z1=z0
|
|
z0=tmp
|
|
|
|
i1=np.argmin(np.abs(s-z0))
|
|
i2=np.argmin(np.abs(s-z1))
|
|
|
|
ylabel=r''
|
|
if filt['BETX']:
|
|
self.plotSingle(s[i1:i2],data['BETX'][i1:i2],(0,0,1,1),r'$\beta_{x}$')
|
|
ylabel=ylabel+r'$\beta_x$ (m), '
|
|
if filt['BETY']:
|
|
self.plotSingle(s[i1:i2],data['BETY'][i1:i2],(1,0,0,1),r'$\beta_{y}$')
|
|
ylabel=ylabel+r'$\beta_y$ (m), '
|
|
if filt['ALFX']:
|
|
self.plotSingle(s[i1:i2],data['ALFX'][i1:i2],(0,0,1,1),r'$\alpha_{x}$')
|
|
ylabel=ylabel+r'$\alpha_x$ (rad), '
|
|
if filt['ALFY']:
|
|
self.plotSingle(s[i1:i2],data['ALFY'][i1:i2],(1,0,0,1),r'$\alpha_{y}$')
|
|
ylabel=ylabel+r'$\alpha_y$ (rad), '
|
|
if filt['DX']:
|
|
self.plotSingle(s[i1:i2],data['DX'][i1:i2],(0,0,1,1),r'$\eta_{x}$')
|
|
ylabel=ylabel+r'$\eta_x$ (m), '
|
|
if filt['DY']:
|
|
self.plotSingle(s[i1:i2],data['DY'][i1:i2],(1,0,0,1),r'$\eta_{y}$')
|
|
ylabel=ylabel+r'$\eta_y$ (m), '
|
|
if filt['RE56']:
|
|
self.plotSingle(s[i1:i2],data['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,data['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
|
|
|
|
|
|
|
|
|
|
# |