def __init__(self, parent, image_item: ChannelImageItem, blend_view):
QWidget.__init__(self, parent)
self.image_item = image_item
self.blend_view = blend_view
layout = QVBoxLayout(self)
layout.setContentsMargins(0, 0, 0, 0)
layout.setSpacing(0)
groupbox = QGroupBox(image_item.channel.label)
layout.addWidget(groupbox)
groupbox_layout = QVBoxLayout()
groupbox_layout.setContentsMargins(0, 0, 0, 0)
groupbox_layout.setSpacing(0)
groupbox.setLayout(groupbox_layout)
graphics_view = GraphicsView(self)
histogram_view = HistogramView(image_item, blend_view)
graphics_view.setCentralItem(histogram_view)
groupbox_layout.addWidget(graphics_view)
class FanDiagram(QtGui.QWidget):
"""
Produces a fan diagram
alphas/gammas should be the matrices saved from the FanCompiler, of form:
arb | niralpha1 | niralpha2 | niralpha3 | niralpha4 | ...
1st sb | 1sb alpha | 1sb alpha | .
2nd sb | 2sb alpha | 2sb alpha | .
3rd sb | 3sb alpha | 3sb alpha | .
.
.
.
Assumes both alphas/gammas are the same shape
Alternatively, pass nirAlpha and SBs as 1st/2nd args (as 1D arrays) to have it
create the fan without any data
Or, if alphaData and gammaData are strings, assumes they're paths to data files to plot
:param alphas:
:param gammas:
:param kwargs:
:return:
"""
def __init__(self, alphaData, gammaData=None, view=None):
super(FanDiagram, self).__init__()
if gammaData is None and isinstance(alphaData, FanCompiler):
alphaData, gammaData = alphaData.build(withErrors=False)[:2]
self.layout = QtWidgets.QHBoxLayout()
self.histAlpha = HistogramLUTWidget(self)
self.centralView = GraphicsView()
self.histGamma = HistogramLUTWidget(self)
self.histAlpha.setMinimumWidth(150)
self.histGamma.setMinimumWidth(150)
self.layout.addWidget(self.histGamma)
self.layout.addWidget(self.centralView)
self.layout.addWidget(self.histAlpha)
self.layout.setContentsMargins(0,0,0,0)
self.layout.setSpacing(0)
self.setLayout(self.layout)
if view is None:
self.view = ViewBox()
else:
self.view = view
self.centralView.setCentralItem(self.view)
self.view.setAspectLocked(True)
self.view.invertY(True)
if isinstance(alphaData, str) and isinstance(gammaData, str):
alphaData = np.genfromtxt(alphaData, delimiter=',')
gammaData = np.genfromtxt(gammaData, delimiter=',')
if alphaData.ndim == gammaData.ndim == 1:
# Assume you just want it to be created, and will later populate it
nirAlphas = alphaData
sbs = gammaData
alphaData = np.ones((sbs.shape[0] + 1, nirAlphas.shape[0] + 1)) * -1
alphaData[1:, 0] = sbs
alphaData[0, 1:] = nirAlphas
gammas = np.ones((sbs.shape[0] + 1, nirAlphas.shape[0] + 1)) * -1
gammas[1:, 0] = sbs
gammas[0, 1:] = nirAlphas
sbs = alphaData[1:, 0]
maxSB = sbs.max()
nirAlphas = alphaData[0, 1:]
self.alphaItem = PolarImageItem(r=sbs, theta=nirAlphas)
self.alphaItem.setImage(alphaData[1:,1:])
# nirAlphas+180 is what causes the gamma angles to appear on the left side of the
# fan. This seemed easier than doing some sort of coordinate inversion/flipping
# on the plot itself.
self.gammaItem = PolarImageItem(sbs, nirAlphas+180, gammaData[1:,1:])
self.view.addItem(self.alphaItem)
self.view.addItem(self.gammaItem)
self.histAlpha.setImageItem(self.alphaItem)
self.histGamma.setImageItem(self.gammaItem)
# manually set the default state to the black-gold-white-green-black. Not sure
# if it's necessary to have this be a free parameter vs being hardcoded
self.histAlpha.gradient.restoreState({
"mode": "rgb",
"ticks": [
(0, (0, 0, 0, 255)),
(.25, (128, 128, 0, 255)),
(.5, (255, 255, 255, 255)),
(.75, (0, 128, 0, 255)),
(1, (0, 0, 0, 255))
]
})
# Set the default spacings for the alpha color axis. Again, not sure if it's
# better to leave the 18pt font hard-coded or not, but I am
self.histAlpha.axis.setTickFont(QtGui.QFont("Arial", 18))
self.histAlpha.axis.setTickSpacing(30, 15)
self.histAlpha.axis.setLabel("α (°)", **{'font-family': 'Times',
"font-size": "18pt"})
# As with alpha, hard code the initial color space for gamma (blue-white-red)
# and the font spacings and stuff
self.histGamma.gradient.restoreState({
"mode": "rgb",
"ticks": [
(0, (255, 0, 0, 255)),
(.5, (255, 255, 255, 255)),
(1, (0, 0, 255, 255))
]
})
self.histGamma.axis.setTickFont(QtGui.QFont("Arial", 18))
self.histGamma.axis.setTickSpacing(15, 15)
self.histGamma.axis.setLabel("γ (°)", **{'font-family': 'Times',
"font-size": "18pt"})
self.histAlpha.item.setLevels(-90, 90)
self.histGamma.item.setLevels(-45, 45)
self.histAlpha.autoHistogramRange()
self.histGamma.autoHistogramRange()
# Make it the right dimensions, making sure that the width is appropriate.
# This makes it easier to automate plotting/saving fans and making sure
# their dimensions are consistent.
g = self.geometry()
# I found these by eye, there's not very much important about them
g.setWidth(773)
g.setHeight(480)
# Manually center it on the screen, since geometry isn't well defined at this point
# before events are processed
g.moveCenter(QtWidgets.QApplication.desktop().screenGeometry().center())
self.setGeometry(g)
# Add in the radial axes for it
self.axes = {
"radial": PolarAxis("radial"),
"azimuthal": PolarAxis("azimuthal")
}
# Lighten the radial font to make it distinct from the other
p = self.axes["radial"].pen()
p.setColor(mkColor("#666666"))
self.axes["radial"].setPen(p)
for a in self.axes.values():
# Make sure the axes sit on top of all other items
a.setZValue(10000)
# make sure that they scale appropriately, instead of just floating on top
a.linkToView(self.view)
# Ignore bounds prevents the window from resizing to try and fit in
# the axes items
self.addItem(a, ignoreBounds=True)
# manually set the positions and string values for alpha angles. [-90, 90] work
# well. The other half needs the +-180 to make sure the gamma angles have the
# correctly labeled with respect to alpha_nir
self.axes["azimuthal"].setTicks(
[
[(ii, str(ii)) for ii in np.arange(-90, 91, 30)] + # alpha side (Q1+Q4)
[(ii, str(ii + 180)) for ii in np.arange(-180, -91, 30)] + #Q3
[(ii, str(ii - 180)) for ii in np.arange(120, 151, 30)], #Q1
]
)
# add a title (without text)
self.titleItem = TextItem()
self.titleItem.setAnchor(Point(0.5, 1)) # anchor on bottom-center
# Again, not sure if it's necessary to have the font color/size being
# a free parameter
self.titleItem.setColor("k")
self.titleItem.setFont(QtGui.QFont("Arial", 15))
# Ignore bounds so that the view won't try to account for it (which
# causes a conflict because the title is placed with respect to the
# view region)
self.view.addItem(self.titleItem, ignoreBounds=True)
self.show()
# Arbitrary forcing updates to try and track down why some things don't
# update correctly
QtWidgets.QApplication.processEvents()
self.view.updateViewRange(True, True)
def setAlphaImage(self, img):
self.alphaItem.setImage(img)
def setGammaImage(self, img):
self.gammaItem.setImage(img)
def setImages(self, alpha, gamma):
self.setAlphaImage(alpha)
self.setGammaImage(gamma)
def export(self, fname, hideHistograms=True,
pngScale = 4):
"""
Save fan diagrams to file, with the full image, and color bars on the alpha/gamma
values
:param fname: the fname to save as
hideHistograms - (True) Prevent rendering the histograms, often ncier for
figures/presentations
If fname.endswith(".svg"), it outputs as an SVG. Howver, it's not the cleanest
thing (the files are quite large/unoptimized, and I can't think of an
easy way to correct that). Also, when the svg is converted to pdf via
Inkscape, things get f****d up for some reasons (axes get thicker, fonts
get borked, pixels get messed up). So, it kinda works, but there's
stuff seriously wrong.
One thing to make things cleaner is to use this site:
https://jakearchibald.github.io/svgomg/
which optimizies the svg and makes it a lot easier to work with
:return:
"""
# defaults = {
# "hideHistograms": False
# }
#
# defaults.update(kwargs)
doSvg = fname.endswith(".svg")
if hideHistograms:
# Hide the histogram data (and shrink the plot)
# to avoid confusing people
self.histAlpha.plot.hide()
self.histAlpha.vb.setMaximumWidth(20)
self.histGamma.plot.hide()
self.histGamma.vb.setMaximumWidth(20)
QtWidgets.QApplication.processEvents()
self.histGamma.axis.setRange(-46.75, 46.75)
self.histAlpha.axis.setRange(-94, 94)
width, height = self.width(), self.height()
if doSvg:
from PyQt5 import QtSvg
outputImage = QtSvg.QSvgGenerator()
outputImage.setFileName(fname)
outputImage.setSize(QtCore.QSize(int(width), int(height)))
# I'm not sure why it has to be this, but the axis on the histogrm
# were fuckingup without it
outputImage.setResolution(96)
else:
outputImage = QtGui.QImage(width * pngScale, height * pngScale,
QtGui.QImage.Format_ARGB32)
# outputImage.setDotsPerMeterX(650 * 100 / 2.54)
# outputImage.setDotsPerMeterY(650 * 100 / 2.54)
# this gives a moderatly high quality image
outputImage.setDevicePixelRatio(pngScale)
outputImage.fill(QtGui.QColor("white"))
outputPainter = QtGui.QPainter(outputImage)
self.render(outputPainter)
if not doSvg:
ret = outputImage.save(fname)
outputPainter.end()
def addItem(self, item, ignoreBounds=False):
self.view.addItem(item, ignoreBounds)
def setViewRadius(self, r):
# Set the view range of the fan diagram such that radius r is visible
self.view.setRange(QtCore.QRect(-r, -r, 2*r, 2*r), padding=0)
def hideHistogramAxes(self, hideTicks=True):
# Hide the histogram region item and plots and all that for
# less cluttered plots. Definitely useful if export is called with
# hideHistograms=True, where the regions are useless.
# Hide the linear regions
self.histGamma.region.hide()
self.histAlpha.region.hide()
# Keep a reference to the old paint methods so you can reverse it if desired
# This stops the painting of the bars which go from the linear region to the
# gradient editor
self.histGamma.item.oldPaint = self.histGamma.item.paint
self.histAlpha.item.oldPaint = self.histAlpha.item.paint
# Overwriting the functions to return None causes all the other rendering
# things to abort
self.histGamma.item.paint = lambda *x: None
self.histAlpha.item.paint = lambda *x: None
if hideTicks:
# Hide the ticks which can be used for changing the stops/colors of
# the gradients, which are rather ugly
# Note: Since this only hides ticks which are present, I don't think
[ii.hide() for ii in self.histAlpha.item.gradient.ticks.keys()]
[ii.hide() for ii in self.histGamma.item.gradient.ticks.keys()]
QtWidgets.QApplication.processEvents()
# Hard coded numbers which make it look like the axes values line up with
# the gradient item, which is more in-line with how color bars are interpreted
self.histGamma.axis.setRange(-46.75, 46.75)
self.histAlpha.axis.setRange(-94, 94)
def showHistogramAxes(self, showTicks=True):
try:
self.histGamma.item.paint = self.histGamma.item.oldPaint
self.histAlpha.item.paint = self.histAlpha.item.oldPaint
del self.histAlpha.item.oldPaint
del self.histGamma.item.oldPaint
except AttributeError:
# You didn't hide them first (or at least not here
return
self.histGamma.region.show()
self.histAlpha.region.show()
if showTicks:
[ii.show() for ii in self.histAlpha.item.gradient.ticks.keys()]
[ii.show() for ii in self.histGamma.item.gradient.ticks.keys()]
@staticmethod
def fromTMatrix(tMatrix, angle = 45, sbs=None):
"""
Create a fan diagram from T matrices directly. The angle needs to be specified
so the T matrix can be converted to a J matrix. The angle is relative to what's
specified in the qwp.extractMatrices.makeU function.
if you pass a string, it assumes it's a file name from a saved one. It'll load
that and plot it. If you also pass values to sbs, it'll make sure only the passed
values are plotted. Otherwise, it'll plot all the sbs in the file
If you pass a tMatrix as returned from the fitting routines, you also need
to pass the sbs directly in this case, since the tMatrices don't include them.
:param tMatrix:
:param angle:
:param sbs:
:return:
"""
if isinstance(tMatrix, str):
# a file is passed
if sbs is not None:
# Pass an array of sbs with a string, and this'll parse
# out the sidebands which aren't included in the passed array
wantsbs = sbs
else:
wantsbs = None
tMatrix, sbs = loadT(tMatrix)
# Handle if only a select number of sidebands is specified
if wantsbs is not None:
try:
# Find the indices of the desired sidebands within the array of
# sidebands actually loaded
wantIdx = [sbs.tolist().index(ii) for ii in wantsbs]
# Cull it to only those specified
sbs = sbs[wantIdx]
# tMatrix is multidimensional (tMatrix.ndim>2), so ellipses cut
# out the other axes
tMatrix = tMatrix[..., wantIdx]
# Ensure that you got everything you want. Could happen if sidebands
# are requested (passed to the function) and not found
assert np.all(wantsbs == sbs)
except ValueError as e:
raise IndexError("Invalid sideband requested ({} is not in loaded)".format(
e.args[0].split(' ')[0]
))
except AssertionError:
raise IndexError("Invalid sideband requested")
jMatrix = makeJfromT(tMatrix, angle)
if sbs is None:
raise RuntimeWarning("Desired sidebands to plot should be specified as kwarg sbs")
sbs = np.arange(8, 38, 2)
alpha, gamma = jonesToFans(sbs = sbs, J=jMatrix)
return FanDiagram(alpha, gamma)
def setTitle(self, title="", adjustBounds=True):
"""
Sets the title of the fan diagram, positioning the text right above the center
of the fan
:param title:
:param adjustBounds:
:return:
"""
self.titleItem.setText(title)
# Move the title so the bottom is at the top of the outer axis
self.titleItem.setPos(0, self.axes["azimuthal"].fullBoundingRect.top())
QtWidgets.QApplication.processEvents()
# Double up because of some weird f*****g issue with Qt not appropriately
# updating things when requested
self.titleItem.setPos(0, self.axes["azimuthal"].fullBoundingRect.top())
QtWidgets.QApplication.processEvents()
# print(self.titleItem.mapRectToView(self.titleItem.boundingRect()))
if adjustBounds:
# Readjust the viewbox to frame the fan better
# Find the top, based on the coordinates of the top of the title
top = self.titleItem.mapRectToView(self.titleItem.boundingRect()).top()
# Bottom is defiend by the bottom of the axes (includes the text)
# Note: this assumes the
bottom = self.axes["azimuthal"].fullBoundingRect.bottom()
# print("bottom", bottom)
w = abs(top-bottom)
# print("new rect", QtCore.QRectF(-w/2, top, w, w))
self.view.setRange(QtCore.QRectF(-w/2, top, w, w), padding=0)
self.view.setRange(QtCore.QRectF(-w/2, top, w, w), padding=0)
# self.view.update()
def setMaxRadius(self, radius=40):
# Set the maximum value for both of the axes to the value specified.
# The 1e-6 is to prevent it from producing an "r=0" label and stuff
self.axes["azimuthal"]._bounds["radial"] = [1e-6, radius]
self.axes["radial"]._bounds["radial"] = [1e-6, radius]
# Need to invalidate the cache for the axes, forcing it to redraw and update
# the bounding rect and stuff
self.axes["azimuthal"].picture = None
self.axes["radial"].picture = None
class astraPlotWidget(QWidget):
twissplotLayout = [
{
'name': 'sigma_x',
'range': [0, 1],
'scale': 1e3
},
{
'name': 'sigma_y',
'range': [0, 1],
'scale': 1e3
},
{
'name': 'kinetic_energy',
'range': [0, 250],
'scale': 1e-6
},
'next_row',
{
'name': 'sigma_p',
'range': [0, 0.015],
'scale': 1e6
},
{
'name': 'sigma_z',
'range': [0, 0.6],
'scale': 1e3
},
{
'name': 'enx',
'range': [0.5, 1.5],
'scale': 1e6
},
'next_row',
{
'name': 'eny',
'range': [0.5, 1.5],
'scale': 1e6
},
{
'name': 'beta_x',
'range': [0, 150],
'scale': 1
},
{
'name': 'beta_y',
'range': [0, 150],
'scale': 1
},
]
def __init__(self, directory='.', **kwargs):
super(astraPlotWidget, self).__init__(**kwargs)
self.beam = raf.beam()
self.twiss = rtf.twiss()
self.directory = directory
''' twissPlotWidget '''
self.twissPlotView = GraphicsView(useOpenGL=True)
self.twissPlotWidget = GraphicsLayout()
self.twissPlotView.setCentralItem(self.twissPlotWidget)
self.latticePlotData = imageio.imread('lattice_plot.png')
self.latticePlots = {}
self.twissPlots = {}
i = -1
for entry in self.twissplotLayout:
if entry == 'next_row':
self.twissPlotWidget.nextRow()
else:
i += 1
p = self.twissPlotWidget.addPlot(title=entry['name'])
p.showGrid(x=True, y=True)
vb = p.vb
vb.setYRange(*entry['range'])
latticePlot = ImageItem(self.latticePlotData)
latticePlot.setOpts(axisOrder='row-major')
vb.addItem(latticePlot)
latticePlot.setZValue(-1) # make sure this image is on top
# latticePlot.setOpacity(0.5)
self.twissPlots[entry['name']] = p.plot(
pen=mkPen('b', width=3))
self.latticePlots[p.vb] = latticePlot
p.vb.sigRangeChanged.connect(self.scaleLattice)
''' beamPlotWidget '''
self.beamPlotWidget = QWidget()
self.beamPlotLayout = QVBoxLayout()
self.item = ImageItem()
self.beamPlotWidget.setLayout(self.beamPlotLayout)
self.beamPlotView = ImageView(imageItem=self.item)
self.rainbow = rainbow()
self.item.setLookupTable(self.rainbow)
self.item.setLevels([0, 1])
# self.beamPlotWidgetGraphicsLayout = GraphicsLayout()
# p = self.beamPlotWidgetGraphicsLayout.addPlot(title='beam')
# p.showGrid(x=True, y=True)
# self.beamPlot = p.plot(pen=None, symbol='+')
# self.beamPlotView.setCentralItem(self.beamPlotWidgetGraphicsLayout)
self.beamPlotXAxisCombo = QComboBox()
self.beamPlotXAxisCombo.addItems(
['x', 'y', 'zn', 'cpx', 'cpy', 'BetaGamma'])
self.beamPlotYAxisCombo = QComboBox()
self.beamPlotYAxisCombo.addItems(
['x', 'y', 'zn', 'cpx', 'cpy', 'BetaGamma'])
self.beamPlotNumberBins = QSpinBox()
self.beamPlotNumberBins.setRange(10, 500)
self.beamPlotNumberBins.setSingleStep(10)
self.histogramBins = 100
self.beamPlotNumberBins.setValue(self.histogramBins)
self.beamPlotAxisWidget = QWidget()
self.beamPlotAxisLayout = QHBoxLayout()
self.beamPlotAxisWidget.setLayout(self.beamPlotAxisLayout)
self.beamPlotAxisLayout.addWidget(self.beamPlotXAxisCombo)
self.beamPlotAxisLayout.addWidget(self.beamPlotYAxisCombo)
self.beamPlotAxisLayout.addWidget(self.beamPlotNumberBins)
self.beamPlotXAxisCombo.currentIndexChanged.connect(self.plotDataBeam)
self.beamPlotYAxisCombo.currentIndexChanged.connect(self.plotDataBeam)
self.beamPlotNumberBins.valueChanged.connect(self.plotDataBeam)
# self.beamPlotXAxisCombo.setCurrentIndex(2)
# self.beamPlotYAxisCombo.setCurrentIndex(5)
self.beamPlotLayout.addWidget(self.beamPlotAxisWidget)
self.beamPlotLayout.addWidget(self.beamPlotView)
''' slicePlotWidget '''
self.sliceParams = [
{
'name': 'slice_normalized_horizontal_emittance',
'units': 'm-rad',
'text': 'enx'
},
{
'name': 'slice_normalized_vertical_emittance',
'units': 'm-rad',
'text': 'eny'
},
{
'name': 'slice_peak_current',
'units': 'A',
'text': 'PeakI'
},
{
'name': 'slice_relative_momentum_spread',
'units': '%',
'text': 'sigma-p'
},
]
self.slicePlotWidget = QWidget()
self.slicePlotLayout = QVBoxLayout()
self.slicePlotWidget.setLayout(self.slicePlotLayout)
# self.slicePlotView = GraphicsView(useOpenGL=True)
self.slicePlotWidgetGraphicsLayout = GraphicsLayoutWidget()
# self.slicePlots = {}
self.slicePlotCheckbox = {}
self.curve = {}
self.sliceaxis = {}
self.slicePlotCheckboxWidget = QWidget()
self.slicePlotCheckboxLayout = QVBoxLayout()
self.slicePlotCheckboxWidget.setLayout(self.slicePlotCheckboxLayout)
self.slicePlot = self.slicePlotWidgetGraphicsLayout.addPlot(
title='Slice', row=0, col=50)
self.slicePlot.showAxis('left', False)
self.slicePlot.showGrid(x=True, y=True)
i = -1
colors = ['b', 'r', 'g', 'k']
for param in self.sliceParams:
i += 1
axis = AxisItem("left")
labelStyle = {'color': '#' + colorStr(mkColor(colors[i]))[0:-2]}
axis.setLabel(text=param['text'],
units=param['units'],
**labelStyle)
viewbox = ViewBox()
axis.linkToView(viewbox)
viewbox.setXLink(self.slicePlot.vb)
self.sliceaxis[param['name']] = [axis, viewbox]
self.curve[param['name']] = PlotDataItem(pen=colors[i], symbol='+')
viewbox.addItem(self.curve[param['name']])
col = self.findFirstEmptyColumnInGraphicsLayout()
self.slicePlotWidgetGraphicsLayout.ci.addItem(axis,
row=0,
col=col,
rowspan=1,
colspan=1)
self.slicePlotWidgetGraphicsLayout.ci.addItem(viewbox,
row=0,
col=50)
p.showGrid(x=True, y=True)
# self.slicePlots[param] = self.slicePlot.plot(pen=colors[i], symbol='+')
self.slicePlotCheckbox[param['name']] = QCheckBox(param['text'])
self.slicePlotCheckboxLayout.addWidget(
self.slicePlotCheckbox[param['name']])
self.slicePlotCheckbox[param['name']].stateChanged.connect(
self.plotDataSlice)
# self.slicePlotView.setCentralItem(self.slicePlotWidgetGraphicsLayout)
self.slicePlotSliceWidthWidget = QSpinBox()
self.slicePlotSliceWidthWidget.setMaximum(1000)
self.slicePlotSliceWidthWidget.setValue(100)
self.slicePlotSliceWidthWidget.setSingleStep(10)
self.slicePlotSliceWidthWidget.setSuffix("fs")
self.slicePlotSliceWidthWidget.setSpecialValueText('Automatic')
self.slicePlotAxisWidget = QWidget()
self.slicePlotAxisLayout = QHBoxLayout()
self.slicePlotAxisWidget.setLayout(self.slicePlotAxisLayout)
self.slicePlotAxisLayout.addWidget(self.slicePlotCheckboxWidget)
self.slicePlotAxisLayout.addWidget(self.slicePlotSliceWidthWidget)
# self.slicePlotXAxisCombo.currentIndexChanged.connect(self.plotDataSlice)
self.slicePlotSliceWidthWidget.valueChanged.connect(
self.changeSliceLength)
# self.beamPlotXAxisCombo.setCurrentIndex(2)
# self.beamPlotYAxisCombo.setCurrentIndex(5)
self.slicePlotLayout.addWidget(self.slicePlotAxisWidget)
self.slicePlotLayout.addWidget(self.slicePlotWidgetGraphicsLayout)
self.layout = QVBoxLayout()
self.setLayout(self.layout)
self.tabWidget = QTabWidget()
self.folderButton = QPushButton('Select Directory')
self.folderLineEdit = QLineEdit()
self.folderLineEdit.setReadOnly(True)
self.folderLineEdit.setText(self.directory)
self.reloadButton = QPushButton()
self.reloadButton.setIcon(qApp.style().standardIcon(
QStyle.SP_BrowserReload))
self.folderWidget = QGroupBox()
self.folderLayout = QHBoxLayout()
self.folderLayout.addWidget(self.folderButton)
self.folderLayout.addWidget(self.folderLineEdit)
self.folderLayout.addWidget(self.reloadButton)
self.folderWidget.setLayout(self.folderLayout)
self.folderWidget.setMaximumWidth(800)
self.reloadButton.clicked.connect(
lambda: self.changeDirectory(self.directory))
self.folderButton.clicked.connect(self.changeDirectory)
self.fileSelector = QComboBox()
self.fileSelector.currentIndexChanged.connect(self.updateScreenCombo)
self.screenSelector = QComboBox()
self.screenSelector.currentIndexChanged.connect(self.changeScreen)
self.beamWidget = QGroupBox()
self.beamLayout = QHBoxLayout()
self.beamLayout.addWidget(self.fileSelector)
self.beamLayout.addWidget(self.screenSelector)
self.beamWidget.setLayout(self.beamLayout)
self.beamWidget.setMaximumWidth(800)
self.beamWidget.setVisible(False)
self.folderBeamWidget = QWidget()
self.folderBeamLayout = QHBoxLayout()
self.folderBeamLayout.setAlignment(Qt.AlignLeft)
self.folderBeamWidget.setLayout(self.folderBeamLayout)
self.folderBeamLayout.addWidget(self.folderWidget)
self.folderBeamLayout.addWidget(self.beamWidget)
self.tabWidget.addTab(self.twissPlotView, 'Twiss Plots')
self.tabWidget.addTab(self.beamPlotWidget, 'Beam Plots')
self.tabWidget.addTab(self.slicePlotWidget, 'Slice Beam Plots')
self.tabWidget.currentChanged.connect(self.changeTab)
self.layout.addWidget(self.folderBeamWidget)
self.layout.addWidget(self.tabWidget)
self.plotType = 'Twiss'
self.changeDirectory(self.directory)
def findFirstEmptyColumnInGraphicsLayout(self):
rowsfilled = self.slicePlotWidgetGraphicsLayout.ci.rows.get(0,
{}).keys()
for i in range(49):
if not i in rowsfilled:
return i
def changeTab(self, i):
if self.tabWidget.tabText(i) == 'Beam Plots':
self.plotType = 'Beam'
self.beamWidget.setVisible(True)
elif self.tabWidget.tabText(i) == 'Slice Beam Plots':
self.plotType = 'Slice'
self.beamWidget.setVisible(True)
else:
self.plotType = 'Twiss'
self.beamWidget.setVisible(False)
self.loadDataFile()
def changeDirectory(self, directory=None):
if directory == None or directory == False:
self.directory = str(
QFileDialog.getExistingDirectory(self, "Select Directory",
self.directory,
QFileDialog.ShowDirsOnly))
else:
self.directory = directory
self.folderLineEdit.setText(self.directory)
self.currentFileText = self.fileSelector.currentText()
self.currentScreenText = self.screenSelector.currentText()
self.getScreenFiles()
self.updateFileCombo()
self.updateScreenCombo()
self.loadDataFile()
def getScreenFiles(self):
self.screenpositions = {}
files = glob.glob(self.directory + '/*.????.???')
filenames = [
'.'.join(os.path.basename(f).split('.')[:-2]) for f in files
]
print 'filenames = ', filenames
runnumber = [os.path.basename(f).split('.')[-1] for f in files]
for f, r in list(set(zip(filenames, runnumber))):
files = glob.glob(self.directory + '/' + f + '.????.???')
screenpositions = [
re.search(f + '\.(\d\d\d\d)\.\d\d\d', s).group(1)
for s in files
]
print 'screenpositions = ', screenpositions
self.screenpositions[f] = {
'screenpositions': sorted(screenpositions),
'run': r
}
def updateFileCombo(self):
self.fileSelector.clear()
i = -1
screenfirstpos = []
for f in self.screenpositions:
screenfirstpos.append(
[f, min(self.screenpositions[f]['screenpositions'])])
screenfirstpos = np.array(screenfirstpos)
sortedscreennames = screenfirstpos[np.argsort(
np.array(screenfirstpos)[:, 1])]
print 'sortedscreennames = ', sortedscreennames
for f in sortedscreennames:
self.fileSelector.addItem(f[0])
i += 1
if f[0] == self.currentFileText:
self.fileSelector.setCurrentIndex(i)
def changeScreen(self, i):
run = self.screenpositions[str(self.fileSelector.currentText())]['run']
self.beamFileName = str(self.fileSelector.currentText()) + '.' + str(
self.screenSelector.currentText()) + '.' + str(run)
# print 'beamFileName = ', self.beamFileName
self.loadDataFile()
def updateScreenCombo(self):
self.screenSelector.clear()
i = -1
for s in self.screenpositions[str(
self.fileSelector.currentText())]['screenpositions']:
self.screenSelector.addItem(s)
i += 1
if s == self.currentScreenText:
self.screenSelector.setCurrentIndex(i)
def loadDataFile(self):
if self.plotType == 'Twiss':
files = sorted(glob.glob(self.directory + "/*Xemit*"))
self.twiss.read_astra_emit_files(files)
self.plotDataTwiss()
elif self.plotType == 'Beam' or self.plotType == 'Slice':
if hasattr(
self,
'beamFileName') and os.path.isfile(self.directory + '/' +
self.beamFileName):
# starttime = time.time()
self.beam.read_astra_beam_file(self.directory + '/' +
self.beamFileName)
# print 'reading file took ', time.time()-starttime, 's'
# print 'Read file: ', self.beamFileName
if self.plotType == 'Beam':
self.plotDataBeam()
else:
self.beam.bin_time()
self.plotDataSlice()
def plotDataTwiss(self):
for entry in self.twissplotLayout:
if entry == 'next_row':
pass
else:
x = self.twiss['z']
y = self.twiss[entry['name']] * entry['scale']
xy = np.transpose(np.array([x, y]))
x, y = np.transpose(xy[np.argsort(xy[:, 0])])
self.twissPlots[entry['name']].setData(x=x,
y=y,
pen=mkPen('b', width=3))
def plotDataBeam(self):
self.histogramBins = self.beamPlotNumberBins.value()
x = getattr(self.beam, str(self.beamPlotXAxisCombo.currentText()))
y = getattr(self.beam, str(self.beamPlotYAxisCombo.currentText()))
h, xedges, yedges = np.histogram2d(x,
y,
self.histogramBins,
normed=True)
x0 = xedges[0]
y0 = yedges[0]
xscale = (xedges[-1] - xedges[0]) / len(xedges)
yscale = (yedges[-1] - yedges[0]) / len(yedges)
self.item.setImage(h)
self.item.setLookupTable(self.rainbow)
# self.item.setLevels([0,1])
def changeSliceLength(self):
self.beam.slice_length = self.slicePlotSliceWidthWidget.value() * 1e-15
self.beam.bin_time()
self.plotDataSlice()
def plotDataSlice(self):
for param in self.sliceParams:
if self.slicePlotCheckbox[param['name']].isChecked():
x = self.beam.slice_bins
self.slicePlot.setRange(xRange=[min(x), max(x)])
# self.plot.setRange(xRange=[-0.5,1.5])
y = getattr(self.beam, param['name'])
self.curve[param['name']].setData(x=x, y=y)
self.sliceaxis[param['name']][0].setVisible(True)
# currentrange = self.sliceaxis[param['name']][0].range
# print 'currentrange = ', currentrange
# self.sliceaxis[param['name']][0].setRange(0, currentrange[1])
else:
# pass
self.curve[param['name']].setData(x=[], y=[])
self.sliceaxis[param['name']][0].setVisible(False)
self.sliceaxis[param['name']][1].autoRange()
currentrange = self.sliceaxis[param['name']][1].viewRange()
self.sliceaxis[param['name']][1].setYRange(0, currentrange[1][1])
def scaleLattice(self, vb, range):
yrange = range[1]
scaleY = 0.05 * abs(yrange[1] - yrange[0])
rect = QRectF(0, yrange[0] + 2 * scaleY, 49.2778, 4 * scaleY)
self.latticePlots[vb].setRect(rect)
class astraPlotWidget(QWidget):
twissplotLayout = [
{
'name': 'sigma_x',
'range': [0, 1],
'scale': 1e3
},
{
'name': 'sigma_y',
'range': [0, 1],
'scale': 1e3
},
{
'name': 'kinetic_energy',
'range': [0, 250],
'scale': 1e-6
},
'next_row',
{
'name': 'sigma_p',
'range': [0, 0.015],
'scale': 1e6
},
{
'name': 'sigma_z',
'range': [0, 0.6],
'scale': 1e3
},
{
'name': 'enx',
'range': [0.5, 1.5],
'scale': 1e6
},
'next_row',
{
'name': 'eny',
'range': [0.5, 1.5],
'scale': 1e6
},
{
'name': 'beta_x',
'range': [0, 150],
'scale': 1
},
{
'name': 'beta_y',
'range': [0, 150],
'scale': 1
},
]
def __init__(self, directory='.', **kwargs):
super(astraPlotWidget, self).__init__(**kwargs)
self.beam = raf.beam()
self.twiss = rtf.twiss()
self.directory = directory
''' twissPlotWidget '''
self.twissPlotView = GraphicsView(useOpenGL=True)
self.twissPlotWidget = GraphicsLayout()
self.twissPlotView.setCentralItem(self.twissPlotWidget)
self.latticePlotData = imageio.imread(
os.path.dirname(os.path.abspath(__file__)) + '/lattice_plot.png')
self.latticePlots = {}
self.twissPlots = {}
i = -1
for entry in self.twissplotLayout:
if entry == 'next_row':
self.twissPlotWidget.nextRow()
else:
i += 1
p = self.twissPlotWidget.addPlot(title=entry['name'])
p.showGrid(x=True, y=True)
vb = p.vb
vb.setYRange(*entry['range'])
latticePlot = ImageItem(self.latticePlotData)
latticePlot.setOpts(axisOrder='row-major')
vb.addItem(latticePlot)
latticePlot.setZValue(-1) # make sure this image is on top
# latticePlot.setOpacity(0.5)
self.twissPlots[entry['name']] = p.plot(
pen=mkPen('b', width=3))
self.latticePlots[p.vb] = latticePlot
p.vb.sigRangeChanged.connect(self.scaleLattice)
''' beamPlotWidget '''
self.beamPlotWidget = QWidget()
self.beamPlotLayout = QVBoxLayout()
self.beamPlotWidget.setLayout(self.beamPlotLayout)
#
# self.beamPlotChoice =
#
self.beamPlotAxisWidget = QWidget()
self.beamPlotAxisWidget.setMaximumHeight(100)
Form = self.beamPlotAxisWidget
self.beamPlotXAxisDict = OrderedDict()
self.beamPlotXAxisDict['x'] = {'scale': 1e3, 'axis': 'x [mm]'}
self.beamPlotXAxisDict['y'] = {'scale': 1e3, 'axis': 'y [mm]'}
self.beamPlotXAxisDict['z'] = {
'scale': 1e6,
'axis': 'z [micron]',
'norm': True
}
self.beamPlotXAxisDict['t'] = {
'scale': 1e12,
'axis': 't [ps]',
'norm': True
}
self.beamPlotXAxisDict['cpx'] = {'scale': 1e3, 'axis': 'cpx [keV]'}
self.beamPlotXAxisDict['cpy'] = {'scale': 1e3, 'axis': 'cpy [keV]'}
self.beamPlotXAxisDict['BetaGamma'] = {
'scale': 0.511,
'axis': 'cp [MeV]'
}
# Form.setObjectName(("Form"))
# Form.resize(874, 212)
sizePolicy = QSizePolicy(QSizePolicy.Preferred, QSizePolicy.Preferred)
sizePolicy.setHorizontalStretch(0)
sizePolicy.setVerticalStretch(0)
sizePolicy.setHeightForWidth(Form.sizePolicy().hasHeightForWidth())
Form.setSizePolicy(sizePolicy)
self.horizontalLayout = QHBoxLayout(Form)
self.horizontalLayout.setObjectName("horizontalLayout")
self.beamPlotXAxisCombo = QComboBox(Form)
self.beamPlotXAxisCombo.addItems(list(self.beamPlotXAxisDict.keys()))
self.beamPlotXAxisCombo.setCurrentIndex(2)
self.horizontalLayout.addWidget(self.beamPlotXAxisCombo)
self.beamPlotYAxisCombo = QComboBox(Form)
self.beamPlotYAxisCombo.addItems(list(self.beamPlotXAxisDict.keys()))
self.beamPlotYAxisCombo.setCurrentIndex(6)
self.horizontalLayout.addWidget(self.beamPlotYAxisCombo)
self.groupBox = QGroupBox(Form)
self.groupBox.setObjectName("groupBox")
self.formLayout = QFormLayout(self.groupBox)
self.formLayout.setFieldGrowthPolicy(QFormLayout.AllNonFixedFieldsGrow)
self.formLayout.setObjectName("formLayout")
self.chooseSliceWidth = QRadioButton(self.groupBox)
self.chooseSliceWidth.setObjectName(("chooseSliceWidth"))
self.formLayout.setWidget(0, QFormLayout.LabelRole,
self.chooseSliceWidth)
self.sliceWidth = QDoubleSpinBox(self.groupBox)
self.sliceWidth.setDecimals(6)
self.sliceWidth.setObjectName("sliceWidth")
self.formLayout.setWidget(0, QFormLayout.FieldRole, self.sliceWidth)
self.chooseSliceNumber = QRadioButton(self.groupBox)
self.chooseSliceNumber.setChecked(True)
self.chooseSliceNumber.setObjectName("chooseSliceNumber")
self.formLayout.setWidget(1, QFormLayout.LabelRole,
self.chooseSliceNumber)
self.sliceNumber = QSpinBox(self.groupBox)
self.sliceNumber.setObjectName(("sliceNumber"))
self.formLayout.setWidget(1, QFormLayout.FieldRole, self.sliceNumber)
self.horizontalLayout.addWidget(self.groupBox)
self.chooseSliceWidth.setText(_translate("Form", "Slice Width", None))
self.chooseSliceNumber.setText(
_translate("Form", "Number of Slices", None))
self.sliceWidth.setRange(1e-6, 1)
self.sliceWidth.setSingleStep(0.00001)
self.histogramWidth = 0.0005
self.sliceWidth.setValue(self.histogramWidth)
# self.sliceNumber = QSpinBox()
self.sliceNumber.setRange(10, 10000)
self.sliceNumber.setSingleStep(10)
self.histogramBins = 100
self.sliceNumber.setValue(self.histogramBins)
# self.beamPlotXAxisCombo = QComboBox()
# self.beamPlotAxisLayout = QHBoxLayout()
# self.beamPlotAxisWidget.setLayout(self.beamPlotAxisLayout)
# self.beamPlotAxisLayout.addWidget(self.beamPlotXAxisCombo)
# self.beamPlotAxisLayout.addWidget(self.beamPlotYAxisCombo)
# self.beamPlotAxisLayout.addWidget(self.beamPlotNumberBins)
self.beamPlotXAxisCombo.currentIndexChanged.connect(self.plotDataBeam)
self.beamPlotYAxisCombo.currentIndexChanged.connect(self.plotDataBeam)
self.chooseSliceNumber.toggled.connect(self.plotDataBeam)
self.sliceNumber.valueChanged.connect(self.plotDataBeam)
self.sliceWidth.valueChanged.connect(self.plotDataBeam)
# self.beamPlotXAxisCombo.setCurrentIndex(2)
# self.beamPlotYAxisCombo.setCurrentIndex(5)
self.canvasWidget = QWidget()
l = QVBoxLayout(self.canvasWidget)
self.sc = MyStaticMplCanvas(self.canvasWidget,
width=1,
height=1,
dpi=150)
l.addWidget(self.sc)
self.beamPlotLayout.addWidget(self.beamPlotAxisWidget)
self.beamPlotLayout.addWidget(self.canvasWidget)
''' slicePlotWidget '''
self.sliceParams = [
{
'name': 'slice_normalized_horizontal_emittance',
'units': 'm-rad',
'text': 'enx'
},
{
'name': 'slice_normalized_vertical_emittance',
'units': 'm-rad',
'text': 'eny'
},
{
'name': 'slice_peak_current',
'units': 'A',
'text': 'PeakI'
},
{
'name': 'slice_relative_momentum_spread',
'units': '%',
'text': 'sigma-p'
},
{
'name': 'slice_beta_x',
'units': 'm',
'text': 'beta_x'
},
{
'name': 'slice_beta_y',
'units': 'm',
'text': 'beta_y'
},
]
self.slicePlotWidget = QWidget()
self.slicePlotLayout = QVBoxLayout()
self.slicePlotWidget.setLayout(self.slicePlotLayout)
# self.slicePlotView = GraphicsView(useOpenGL=True)
self.slicePlotWidgetGraphicsLayout = GraphicsLayoutWidget()
# self.slicePlots = {}
self.slicePlotCheckbox = {}
self.curve = {}
self.sliceaxis = {}
self.slicePlotCheckboxWidget = QWidget()
self.slicePlotCheckboxLayout = QVBoxLayout()
self.slicePlotCheckboxWidget.setLayout(self.slicePlotCheckboxLayout)
self.slicePlot = self.slicePlotWidgetGraphicsLayout.addPlot(
title='Slice', row=0, col=50)
self.slicePlot.showAxis('left', False)
self.slicePlot.showGrid(x=True, y=True)
i = -1
colors = ['b', 'r', 'g', 'k', 'y', 'm', 'c']
for param in self.sliceParams:
i += 1
axis = AxisItem("left")
labelStyle = {'color': '#' + colorStr(mkColor(colors[i]))[0:-2]}
axis.setLabel(text=param['text'],
units=param['units'],
**labelStyle)
viewbox = ViewBox()
axis.linkToView(viewbox)
viewbox.setXLink(self.slicePlot.vb)
self.sliceaxis[param['name']] = [axis, viewbox]
self.curve[param['name']] = PlotDataItem(pen=colors[i], symbol='+')
viewbox.addItem(self.curve[param['name']])
col = self.findFirstEmptyColumnInGraphicsLayout()
self.slicePlotWidgetGraphicsLayout.ci.addItem(axis,
row=0,
col=col,
rowspan=1,
colspan=1)
self.slicePlotWidgetGraphicsLayout.ci.addItem(viewbox,
row=0,
col=50)
p.showGrid(x=True, y=True)
# self.slicePlots[param] = self.slicePlot.plot(pen=colors[i], symbol='+')
self.slicePlotCheckbox[param['name']] = QCheckBox(param['text'])
self.slicePlotCheckboxLayout.addWidget(
self.slicePlotCheckbox[param['name']])
self.slicePlotCheckbox[param['name']].stateChanged.connect(
self.plotDataSlice)
# self.slicePlotView.setCentralItem(self.slicePlotWidgetGraphicsLayout)
self.slicePlotSliceWidthWidget = QSpinBox()
self.slicePlotSliceWidthWidget.setMaximum(500)
self.slicePlotSliceWidthWidget.setValue(100)
self.slicePlotSliceWidthWidget.setSingleStep(10)
self.slicePlotSliceWidthWidget.setSuffix(" slices")
self.slicePlotSliceWidthWidget.setSpecialValueText('Automatic')
self.slicePlotAxisWidget = QWidget()
self.slicePlotAxisLayout = QHBoxLayout()
self.slicePlotAxisWidget.setLayout(self.slicePlotAxisLayout)
self.slicePlotAxisLayout.addWidget(self.slicePlotCheckboxWidget)
self.slicePlotAxisLayout.addWidget(self.slicePlotSliceWidthWidget)
# self.slicePlotXAxisCombo.currentIndexChanged.connect(self.plotDataSlice)
self.slicePlotSliceWidthWidget.valueChanged.connect(
self.changeSliceLength)
# self.beamPlotXAxisCombo.setCurrentIndex(2)
# self.beamPlotYAxisCombo.setCurrentIndex(5)
self.slicePlotLayout.addWidget(self.slicePlotAxisWidget)
self.slicePlotLayout.addWidget(self.slicePlotWidgetGraphicsLayout)
self.layout = QVBoxLayout()
self.setLayout(self.layout)
self.tabWidget = QTabWidget()
self.folderButton = QPushButton('Select Directory')
self.folderLineEdit = QLineEdit()
self.folderLineEdit.setReadOnly(True)
self.folderLineEdit.setText(self.directory)
self.reloadButton = QPushButton()
self.reloadButton.setIcon(qApp.style().standardIcon(
QStyle.SP_BrowserReload))
self.folderWidget = QGroupBox()
self.folderLayout = QHBoxLayout()
self.folderLayout.addWidget(self.folderButton)
self.folderLayout.addWidget(self.folderLineEdit)
self.folderLayout.addWidget(self.reloadButton)
self.folderWidget.setLayout(self.folderLayout)
self.folderWidget.setMaximumWidth(800)
self.reloadButton.clicked.connect(
lambda: self.changeDirectory(self.directory))
self.folderButton.clicked.connect(self.changeDirectory)
self.fileSelector = QComboBox()
self.fileSelector.currentIndexChanged.connect(self.updateScreenCombo)
self.screenSelector = QComboBox()
self.screenSelector.currentIndexChanged.connect(self.changeScreen)
self.beamWidget = QGroupBox()
self.beamLayout = QHBoxLayout()
self.beamLayout.addWidget(self.fileSelector)
self.beamLayout.addWidget(self.screenSelector)
self.beamWidget.setLayout(self.beamLayout)
self.beamWidget.setMaximumWidth(800)
self.beamWidget.setVisible(False)
self.folderBeamWidget = QWidget()
self.folderBeamLayout = QHBoxLayout()
self.folderBeamLayout.setAlignment(Qt.AlignLeft)
self.folderBeamWidget.setLayout(self.folderBeamLayout)
self.folderBeamLayout.addWidget(self.folderWidget)
self.folderBeamLayout.addWidget(self.beamWidget)
self.tabWidget.addTab(self.twissPlotView, 'Twiss Plots')
self.tabWidget.addTab(self.beamPlotWidget, 'Beam Plots')
self.tabWidget.addTab(self.slicePlotWidget, 'Slice Beam Plots')
# self.log = lw.loggerWidget()
# self.log.addLogger(logger)
# sys.stdout = lw.redirectLogger(self.log, 'stdout')
# self.tabWidget.addTab(self.log,'Log')
self.tabWidget.currentChanged.connect(self.changeTab)
self.layout.addWidget(self.folderBeamWidget)
self.layout.addWidget(self.tabWidget)
self.plotType = 'Twiss'
self.changeDirectory(self.directory)
def findFirstEmptyColumnInGraphicsLayout(self):
rowsfilled = list(
self.slicePlotWidgetGraphicsLayout.ci.rows.get(0, {}).keys())
for i in range(49):
if not i in rowsfilled:
return i
def changeTab(self, i):
if self.tabWidget.tabText(i) == 'Beam Plots':
self.plotType = 'Beam'
self.beamWidget.setVisible(True)
elif self.tabWidget.tabText(i) == 'Slice Beam Plots':
self.plotType = 'Slice'
self.beamWidget.setVisible(True)
else:
self.plotType = 'Twiss'
self.beamWidget.setVisible(False)
self.loadDataFile()
def changeDirectory(self, directory=None):
self.screenSelector.currentIndexChanged.disconnect(self.changeScreen)
self.fileSelector.currentIndexChanged.disconnect(
self.updateScreenCombo)
if directory == None or directory == False:
self.directory = str(
QFileDialog.getExistingDirectory(self, "Select Directory",
self.directory,
QFileDialog.ShowDirsOnly))
else:
self.directory = directory
self.folderLineEdit.setText(self.directory)
self.currentFileText = self.fileSelector.currentText()
self.currentScreenText = self.screenSelector.currentText()
self.getScreenFiles()
self.updateFileCombo()
self.updateScreenCombo()
self.loadDataFile()
self.screenSelector.currentIndexChanged.connect(self.changeScreen)
self.fileSelector.currentIndexChanged.connect(self.updateScreenCombo)
self.changeScreen(self.screenSelector.currentIndex())
def getSposition(self, file):
file = h5py.File(self.directory + '/' + file + '.hdf5', "r")
zpos = file.get('/Parameters/Starting_Position')[2]
if abs(zpos) < 0.01:
print(zpos, file)
return zpos
def getScreenFiles(self):
self.screenpositions = {}
files = glob.glob(self.directory + '/*.hdf5')
filenames = [
'-'.join(os.path.basename(f).split('-')[:2]) for f in files
]
# print 'filenames = ', filenames
runnumber = ['001' for f in filenames]
# print 'runnumber = ', runnumber
for f, r in list(set(zip(filenames, runnumber))):
files = glob.glob(self.directory + '/' + f + '*.hdf5')
screennames = sorted(
[os.path.basename(s).split('.')[0] for s in files],
key=lambda x: self.getSposition(x))
screenpositions = [self.getSposition(s) for s in screennames]
# print 'screenpositions = ', screenpositions
self.screenpositions[f] = {
'screenpositions': screenpositions,
'screennames': screennames,
'run': r
}
def updateFileCombo(self):
self.fileSelector.clear()
i = -1
screenfirstpos = []
for f in self.screenpositions:
if len(self.screenpositions[f]['screenpositions']) > 0:
screenfirstpos.append(
[f, min(self.screenpositions[f]['screenpositions'])])
screenfirstpos = np.array(screenfirstpos)
# print 'screenfirstpos = ', screenfirstpos
sortedscreennames = sorted(screenfirstpos, key=lambda x: float(x[1]))
print('sortedscreennames = ', sortedscreennames)
for f in sortedscreennames:
self.fileSelector.addItem(f[0])
i += 1
if f[0] == self.currentFileText:
self.fileSelector.setCurrentIndex(i)
def changeScreen(self, i):
run = self.screenpositions[str(self.fileSelector.currentText())]['run']
self.beamFileName = str(self.screenpositions[str(
self.fileSelector.currentText())]['screennames'][
self.screenSelector.currentIndex()]) + '.hdf5'
print('beamFileName = ', self.beamFileName)
self.loadDataFile()
def updateScreenCombo(self):
self.screenSelector.clear()
i = -1
for i, s in enumerate(self.screenpositions[str(
self.fileSelector.currentText())]['screennames']):
self.screenSelector.addItem(
str(s) + ' (' + str(self.screenpositions[str(
self.fileSelector.currentText())]['screenpositions'][i]) +
')')
i += 1
if s == self.currentScreenText:
self.screenSelector.setCurrentIndex(i)
def loadDataFile(self):
if self.plotType == 'Twiss':
files = sorted(glob.glob(self.directory + "/*Xemit*"))
self.twiss.read_astra_emit_files(files)
files = sorted(glob.glob(self.directory + "/*.flr"))
self.twiss.read_elegant_twiss_files(files, reset=False)
self.plotDataTwiss()
elif self.plotType == 'Beam' or self.plotType == 'Slice':
if hasattr(
self,
'beamFileName') and os.path.isfile(self.directory + '/' +
self.beamFileName):
# starttime = time.time()
self.beam.read_HDF5_beam_file(self.directory + '/' +
self.beamFileName)
# print 'reading file took ', time.time()-starttime, 's'
# print 'Read file: ', self.beamFileName
if self.plotType == 'Beam':
self.plotDataBeam()
else:
self.changeSliceLength()
# self.plotDataSlice()
def plotDataTwiss(self):
for entry in self.twissplotLayout:
if entry == 'next_row':
pass
else:
x = self.twiss['z']
y = self.twiss[entry['name']] * entry['scale']
xy = np.transpose(np.array([x, y]))
x, y = np.transpose(xy[np.argsort(xy[:, 0])])
self.twissPlots[entry['name']].setData(x=x,
y=y,
pen=mkPen('b', width=3))
def plotDataBeam(self):
xdict = self.beamPlotXAxisDict[str(
self.beamPlotXAxisCombo.currentText())]
ydict = self.beamPlotXAxisDict[str(
self.beamPlotYAxisCombo.currentText())]
x = xdict['scale'] * getattr(
self.beam, str(self.beamPlotXAxisCombo.currentText()))
if 'norm' in xdict and xdict['norm'] == True:
x = x - np.mean(x)
y = ydict['scale'] * getattr(
self.beam, str(self.beamPlotYAxisCombo.currentText()))
if 'norm' in ydict and ydict['norm'] == True:
y = y - np.mean(y)
if self.chooseSliceWidth.isChecked():
slices = np.arange(min(x),
max(x) + self.sliceWidth.value(),
self.sliceWidth.value())
self.sliceNumber.setValue(len(slices))
self.sc.plothist(x,
y, [slices, len(slices)],
xLabel=xdict['axis'],
yLabel=ydict['axis'])
else:
self.sliceWidth.setValue(
int(1e6 * ((max(x) - min(x)) / self.sliceNumber.value())) /
1e6)
self.sc.plothist(x,
y,
self.sliceNumber.value(),
xLabel=xdict['axis'],
yLabel=ydict['axis'])
def changeSliceLength(self):
self.beam.slices = self.slicePlotSliceWidthWidget.value()
self.beam.bin_time()
self.plotDataSlice()
def plotDataSlice(self):
for param in self.sliceParams:
if self.slicePlotCheckbox[param['name']].isChecked():
exponent = np.floor(np.log10(np.abs(self.beam.slice_length)))
x = 10**(12) * np.array(
(self.beam.slice_bins - np.mean(self.beam.slice_bins)))
self.slicePlot.setRange(xRange=[min(x), max(x)])
# self.plot.setRange(xRange=[-0.5,1.5])
y = getattr(self.beam, param['name'])
self.curve[param['name']].setData(x=x, y=y)
self.sliceaxis[param['name']][0].setVisible(True)
# currentrange = self.sliceaxis[param['name']][0].range
# print 'currentrange = ', currentrange
# self.sliceaxis[param['name']][0].setRange(0, currentrange[1])
else:
# pass
self.curve[param['name']].setData(x=[], y=[])
self.sliceaxis[param['name']][0].setVisible(False)
self.sliceaxis[param['name']][1].autoRange()
currentrange = self.sliceaxis[param['name']][1].viewRange()
self.sliceaxis[param['name']][1].setYRange(0, currentrange[1][1])
def scaleLattice(self, vb, range):
yrange = range[1]
scaleY = 0.05 * abs(yrange[1] - yrange[0])
rect = QRectF(0, yrange[0] + 2 * scaleY, 49.2778, 4 * scaleY)
self.latticePlots[vb].setRect(rect)
