예제 #1
0
    def __init__(self, parent: QWidget):
        GraphicsView.__init__(self, parent)

        self.layout = GraphicsLayout()
        self.setCentralItem(self.layout)

        self.tiles: Dict[str, TileView] = dict()
예제 #2
0
    def __init__(self):
        QtGui.QWidget.__init__(self)

        self.data = {}

        self.gv = pg.GraphicsView()
        self.gl = GraphicsLayout()

        self.gv.setCentralItem(self.gl)

        self.layout = QtGui.QGridLayout()
        self.layout.addWidget(self.gv, 1, 1, 1, 1)

        self.setLayout(self.layout)
예제 #3
0
    def __init__(self, title):
        self._app = QtGui.QApplication([])
        super().__init__(
        )  # GraphicsView requires a QApplication before it can be made
        super().setWindowTitle(title)

        self.layout = GraphicsLayout()  # for ease of organization
        super().setCentralItem(self.layout)
예제 #4
0
    def __init__(self, parent: QWidget):
        GraphicsView.__init__(self, parent)
        layout = GraphicsLayout()
        self.setCentralItem(layout)

        self._image_item = ImageItem()
        self.viewbox = ViewBox(layout, lockAspect=True, invertY=True)
        self.viewbox.addItem(self._image_item)
        layout.addItem(self.viewbox)

        self.scale = ScaleBar(size=10, suffix='μm')
        self.scale.setParentItem(self.viewbox)
        self.scale.anchor((1, 1), (1, 1), offset=(-20, -20))
        self.scale.hide()

        self._show_mask = False
        self.blend_mode = QPainter.CompositionMode_Screen
        self.items: List[ChannelImageItem] = None
예제 #5
0
    def __init__(self):
        super().__init__()

        self._layout = GraphicsLayout()  # 创建一个网格布局
        self._layout.setContentsMargins(10, 10, 10, 10)
        self.setCentralItem(self._layout)

        label = pg.LabelItem()

        label.setText("这是一个label Item")
        label.item.setPos(100, 100)

        text = pg.TextItem()
        text.setText("这是一个text Item")

        plot_01 = pg.PlotItem()
        plot_01.addItem(label)
        plot_01.setMenuEnabled(False)
        plot_01.setClipToView(True)
        plot_01.hideAxis('left')  # 隐藏左纵坐标轴
        plot_01.hideAxis('bottom')

        plot_01.showAxis('right')  # 显示右纵坐标轴
        # 字体沿Y轴翻转
        plot_01.invertY()
        plot_01.setDownsampling(mode='peak')

        plot_01.setRange(xRange=(0, 1), yRange=(0, 1))

        axis_width = plot_01.getAxis("right").width()
        axis_height = plot_01.getAxis("bottom").height()
        axis_offset = QtCore.QPointF(axis_width, axis_height)
        bottom_view = plot_01.getViewBox()
        # bottom_right= bottom_view.mapSceneToView(
        # bottom_view.sceneBoundingRect().bottomRight() - axis_offset)
        # top_left = bottom_view.mapSceneToView(bottom_view.sceneBoundingRect().topLeft())
        # label.setPos(top_left)
        # plot_02 = pg.PlotItem()
        # plot_02.addItem(label)

        #
        self._layout.addItem(plot_01)
        self._layout.nextRow()
        self._layout.addItem(plot_01)
예제 #6
0
class TilesView(GraphicsView):
    def __init__(self, parent: QWidget):
        GraphicsView.__init__(self, parent)

        self.layout = GraphicsLayout()
        self.setCentralItem(self.layout)

        self.tiles: Dict[str, TileView] = dict()

    def clear(self):
        self.layout.clear()
        self.tiles.clear()

    def get_cell(self, i: int):
        l = len(self.tiles.keys())
        if l > 9:
            rows = (0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3)
            cols = (0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3)
        elif l > 4:
            rows = (0, 0, 0, 1, 1, 1, 2, 2, 2)
            cols = (0, 1, 2, 0, 1, 2, 0, 1, 2)
        else:
            rows = (0, 0, 1, 1)
            cols = (0, 1, 0, 1)
        return rows[i], cols[i]

    def set_images(self, items: List[ChannelImageItem]):
        self.clear()
        if len(items) > 0:
            for item in items:
                tile = TileView(self.layout, item)
                self.tiles[item.channel.metal] = tile

            first_tile = self.tiles[list(self.tiles.keys())[0]]
            for i, tile in enumerate(self.tiles.values()):
                if first_tile is not tile:
                    tile.linkView(ViewBox.XAxis, first_tile)
                    tile.linkView(ViewBox.YAxis, first_tile)
                cell = self.get_cell(i)
                self.layout.addItem(tile, cell[0], cell[1])

    def fit_all_tiles(self):
        for name, tile in self.tiles.items():
            tile.autoRange()
예제 #7
0
    def __init__(self, interType, ui):  #Construtor da classe
        self.ui = ui
        self.interType = interType

        self.timer = QtCore.QTimer()

        self.ser1 = SerialManager(interType)

        self.ySize = 500

        if self.interType == 0:
            self.x_scale = 10
            self.nCurves = 2
            self.y_min = [0, 0]
            self.y_max = [32700, 10000]
            self.xSize = 780

            self.color = ['r', 'g']
            self.name = ["Força", "Tensão no Calibrante"]
            self.unit = ["Tonf", "mV"]

        elif self.interType == 1:
            self.x_scale = 10
            self.nCurves = 3
            self.y_min = [0, 0, 0]
            self.y_max = [32700, 10000, 5000]

            self.color = ['r', 'g', 'y']
            self.name = ["Temperatura", "Força", "Potência"]
            self.unit = ["ºC", "N", "W"]
            self.xSize = 735

        self.ser1.startPort(str(self.ser1.portList[1].device), 115200)
        self.ser1.serialListPanel(self.ui)

        self.layout = GraphicsLayout()

        self.updateTimer()

        self.scene = QtGui.QGraphicsScene()

        self.scene.addItem(self.layout)
        self.updateScale()

        #self.layout.setMaximumSize(770,550)
        self.scene.focusItem()

        self.sceneSelector(self.scene)

        self.ui.menuPlay_Pause = QtWidgets.QAction(self.ui.MainWindow)
        self.ui.menuPlay_Pause.setText("Play")
        self.ui.menuPlay_Pause.setObjectName("menuPlay_Pause")

        self.ui.menuBar.addAction(self.ui.menuPlay_Pause)
        self.ui.menuPlay_Pause.setCheckable(True)

        self.ui.menuPlay_Pause.triggered.connect(self.playPauseButtonAnimation)

        self.ui.linkActions()

        self.ui.t_max.returnPressed.connect(self.updateScale)
        self.ui.f_max.returnPressed.connect(self.updateScale)
        self.ui.f_min.returnPressed.connect(self.updateScale)
        self.ui.p_max.returnPressed.connect(self.updateScale)
        self.ui.p_min.returnPressed.connect(self.updateScale)

        if self.interType == 1:
            self.ui.temp_max.returnPressed.connect(self.updateScale)
            self.ui.temp_min.returnPressed.connect(self.updateScale)

        self.ui.samplingCBox.currentIndexChanged.connect(self.updateTimer)
        self.timeant = time.time()
        self.timer.timeout.connect(self.updateData)

        thread_instance = QtCore.QThread()
        thread_instance.start()
        thread_instance.exec_()
        self.timer.start(100)
예제 #8
0
class SystemEngine(object):
    def __init__(self, interType, ui):  #Construtor da classe
        self.ui = ui
        self.interType = interType

        self.timer = QtCore.QTimer()

        self.ser1 = SerialManager(interType)

        self.ySize = 500

        if self.interType == 0:
            self.x_scale = 10
            self.nCurves = 2
            self.y_min = [0, 0]
            self.y_max = [32700, 10000]
            self.xSize = 780

            self.color = ['r', 'g']
            self.name = ["Força", "Tensão no Calibrante"]
            self.unit = ["Tonf", "mV"]

        elif self.interType == 1:
            self.x_scale = 10
            self.nCurves = 3
            self.y_min = [0, 0, 0]
            self.y_max = [32700, 10000, 5000]

            self.color = ['r', 'g', 'y']
            self.name = ["Temperatura", "Força", "Potência"]
            self.unit = ["ºC", "N", "W"]
            self.xSize = 735

        self.ser1.startPort(str(self.ser1.portList[1].device), 115200)
        self.ser1.serialListPanel(self.ui)

        self.layout = GraphicsLayout()

        self.updateTimer()

        self.scene = QtGui.QGraphicsScene()

        self.scene.addItem(self.layout)
        self.updateScale()

        #self.layout.setMaximumSize(770,550)
        self.scene.focusItem()

        self.sceneSelector(self.scene)

        self.ui.menuPlay_Pause = QtWidgets.QAction(self.ui.MainWindow)
        self.ui.menuPlay_Pause.setText("Play")
        self.ui.menuPlay_Pause.setObjectName("menuPlay_Pause")

        self.ui.menuBar.addAction(self.ui.menuPlay_Pause)
        self.ui.menuPlay_Pause.setCheckable(True)

        self.ui.menuPlay_Pause.triggered.connect(self.playPauseButtonAnimation)

        self.ui.linkActions()

        self.ui.t_max.returnPressed.connect(self.updateScale)
        self.ui.f_max.returnPressed.connect(self.updateScale)
        self.ui.f_min.returnPressed.connect(self.updateScale)
        self.ui.p_max.returnPressed.connect(self.updateScale)
        self.ui.p_min.returnPressed.connect(self.updateScale)

        if self.interType == 1:
            self.ui.temp_max.returnPressed.connect(self.updateScale)
            self.ui.temp_min.returnPressed.connect(self.updateScale)

        self.ui.samplingCBox.currentIndexChanged.connect(self.updateTimer)
        self.timeant = time.time()
        self.timer.timeout.connect(self.updateData)

        thread_instance = QtCore.QThread()
        thread_instance.start()
        thread_instance.exec_()
        self.timer.start(100)

    def playPauseButtonAnimation(self):
        if (self.ui.menuPlay_Pause.text() == "Play"):
            self.ui.menuPlay_Pause.setText("Pause")
            self.ui.startTimeLabel.setText(self.ui.currentTimeLabel.text())
        else:
            self.ui.menuPlay_Pause.setText("Play")

    def sceneSelector(self, scene):
        scene.setSceneRect(0, 0, 1000, 600)
        self.ui.CentralGraph.setScene(scene)
        self.ui.CentralGraph.setSceneRect(scene.sceneRect())
        self.ui.CentralGraph.setBackgroundBrush(QtCore.Qt.black)
        self.ui.CentralGraph.setInteractive(False)

    def updateData(self):
        #print(time.time()-self.timeant)
        if (self.ui.menuPlay_Pause.isChecked() == True):
            readData = self.ser1.read()  # Lê o dado da serial
            readData = readData.decode('utf8')
            #print(readData)
            dado = readData.split(' ', self.nCurves + 1)

            dado[len(dado) - 1] = dado[len(dado) - 1].split('\n', 2)[0]
            try:
                if self.interType == 0:
                    self.ui.forceLabel.setText(str(dado[0]) + " Tonf")
                    self.ui.calibratorLabel.setText(
                        str(dado[1].split()[0]) + " mV")

                elif self.interType == 1:
                    self.ui.tempLabel.setText(str(dado[0]) + "ºC")
                    self.ui.forceLabel.setText(str(dado[1].split()[0]) + " kN")
                    self.ui.powLabel.setText(str(dado[2].split()[0]) + " kW")

            except IndexError:
                print("Erro no Indice do Array Enviado pela Serial")
            try:
                self.graph.updateGraph(dado)

            except IndexError:
                print("Erro: Array inválido")
                print(dado)

        else:
            pass
        self.timeant = time.time()

    def updateScale(self):
        try:
            self.x_scale = int(self.ui.t_max.text())

            if self.interType == 0:
                self.y_min[0] = int(self.ui.f_min.text())
                self.y_max[0] = int(self.ui.f_max.text())
                self.y_min[1] = int(self.ui.p_min.text())
                self.y_max[1] = int(self.ui.p_max.text())

            elif self.interType == 1:
                self.y_min[0] = int(self.ui.temp_min.text())
                self.y_max[0] = int(self.ui.temp_max.text())
                self.y_min[1] = int(self.ui.f_min.text())
                self.y_max[1] = int(self.ui.f_max.text())
                self.y_max[2] = int(self.ui.p_max.text())
                self.y_min[2] = int(self.ui.p_min.text())

        except ValueError:
            self.x_scale = 100
            self.y_min[0] = 0
            self.y_max[0] = 32700
            self.y_min[1] = 0
            self.y_max[1] = 10000

            if self.interType == 1:
                self.y_min[2] = 0
                self.y_max[2] = 5000

            print("Erro!: Campos de Escala Vazios")

        if (self.x_scale == "0"):
            self.x_scale = "1"

        self.layout.clear()
        self.layout.setMinimumSize(500, 600)
        self.layout.setMaximumSize(500, 600)
        self.graph = Graph(self.time[0] / 1000, self.x_scale, self.nCurves,
                           self.y_min, self.y_max, self.xSize, self.ySize,
                           self.color, self.name, self.unit)

        self.layout.addItem(self.graph.axis[0],
                            row=1,
                            col=1,
                            rowspan=1,
                            colspan=1)
        for n in range(1, self.nCurves):
            self.layout.addItem(self.graph.axis[n],
                                row=1,
                                col=n + 2,
                                rowspan=1,
                                colspan=1)

        self.layout.addItem(self.graph, row=1, col=2, rowspan=1, colspan=1)
        self.layout.addItem(self.graph.axisTime,
                            row=2,
                            col=2,
                            rowspan=1,
                            colspan=1)

        self.scene.focusItem()

    def updateTimer(self):
        self.time = self.ui.samplingCBox.currentText().split(" ")
        self.time[0] = int(self.time[0])
        #print(self.time[0])
        if (self.time[0] < 100):
            self.time[0] = 1000 * self.time[0]

        self.timer.stop()
        self.timer.start(self.time[0])
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)
    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 __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)
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)
    def __init__(self, parent, shared_data):
        super(DataWidget, self).__init__(parent)
        self.__shared_data = shared_data
        self.__shared_data.update_sync.emit()

        # Add the file selection controls
        self.__dir_picker_button = QPushButton()
        self.__dir_picker_button.setEnabled(True)
        self.__dir_picker_button.setText("Load data")
        self.__dir_picker_button.setIcon(self.style().standardIcon(QStyle.SP_DirIcon))
        self.__dir_picker_button.setToolTip('Select the directory using the file explorer')
        self.__dir_picker_button.clicked.connect(self.__open_dir_picker)

        # Add the sync controls
        self.__sync_time_label = QLabel()
        self.__sync_time_label.setText('Enter the timecode (HH:mm:ss:zzz) : ')

        self.__sync_time_edit = QTimeEdit()
        self.__sync_time_edit.setDisplayFormat('HH:mm:ss:zzz')
        self.__sync_time_edit.setEnabled(False)

        self.__sync_time_button = QPushButton()
        self.__sync_time_button.setText('Sync data')
        self.__sync_time_button.setEnabled(False)
        self.__sync_time_button.clicked.connect(self.__sync_data)

        # Create the layout for the file controls
        dir_layout = QHBoxLayout()
        dir_layout.setContentsMargins(0, 0, 0, 0)
        dir_layout.addWidget(self.__dir_picker_button)
        dir_layout.addStretch(1)
        dir_layout.addWidget(self.__sync_time_label)
        dir_layout.addWidget(self.__sync_time_edit)
        dir_layout.addWidget(self.__sync_time_button)

        # Create the axis and their viewbox
        self.__x_axis_item = AxisItem('left')
        self.__y_axis_item = AxisItem('left')
        self.__z_axis_item = AxisItem('left')

        self.__x_axis_viewbox = ViewBox()
        self.__y_axis_viewbox = ViewBox()
        self.__z_axis_viewbox = ViewBox()

        # Create the widget which will display the data
        self.__graphic_view = GraphicsView(background="#ecf0f1")
        self.__graphic_layout = GraphicsLayout()
        self.__graphic_view.setCentralWidget(self.__graphic_layout)

        # Add the axis to the widget
        self.__graphic_layout.addItem(self.__x_axis_item, row=2, col=3, rowspan=1, colspan=1)
        self.__graphic_layout.addItem(self.__y_axis_item, row=2, col=2, rowspan=1, colspan=1)
        self.__graphic_layout.addItem(self.__z_axis_item, row=2, col=1, rowspan=1, colspan=1)

        self.__plot_item = PlotItem()
        self.__plot_item_viewbox = self.__plot_item.vb
        self.__graphic_layout.addItem(self.__plot_item, row=2, col=4, rowspan=1, colspan=1)

        self.__graphic_layout.scene().addItem(self.__x_axis_viewbox)
        self.__graphic_layout.scene().addItem(self.__y_axis_viewbox)
        self.__graphic_layout.scene().addItem(self.__z_axis_viewbox)

        self.__x_axis_item.linkToView(self.__x_axis_viewbox)
        self.__y_axis_item.linkToView(self.__y_axis_viewbox)
        self.__z_axis_item.linkToView(self.__z_axis_viewbox)

        self.__x_axis_viewbox.setXLink(self.__plot_item_viewbox)
        self.__y_axis_viewbox.setXLink(self.__plot_item_viewbox)
        self.__z_axis_viewbox.setXLink(self.__plot_item_viewbox)

        self.__plot_item_viewbox.sigResized.connect(self.__update_views)
        self.__x_axis_viewbox.enableAutoRange(axis=ViewBox.XAxis, enable=True)
        self.__y_axis_viewbox.enableAutoRange(axis=ViewBox.XAxis, enable=True)
        self.__z_axis_viewbox.enableAutoRange(axis=ViewBox.XAxis, enable=True)

        # Create the final layout
        self.__v_box = QVBoxLayout()
        self.__v_box.addLayout(dir_layout)
        self.__v_box.addWidget(self.__graphic_view)

        self.setLayout(self.__v_box)

        self.__restore_state()
class DataWidget(QWidget):
    def __init__(self, parent, shared_data):
        super(DataWidget, self).__init__(parent)
        self.__shared_data = shared_data
        self.__shared_data.update_sync.emit()

        # Add the file selection controls
        self.__dir_picker_button = QPushButton()
        self.__dir_picker_button.setEnabled(True)
        self.__dir_picker_button.setText("Load data")
        self.__dir_picker_button.setIcon(self.style().standardIcon(QStyle.SP_DirIcon))
        self.__dir_picker_button.setToolTip('Select the directory using the file explorer')
        self.__dir_picker_button.clicked.connect(self.__open_dir_picker)

        # Add the sync controls
        self.__sync_time_label = QLabel()
        self.__sync_time_label.setText('Enter the timecode (HH:mm:ss:zzz) : ')

        self.__sync_time_edit = QTimeEdit()
        self.__sync_time_edit.setDisplayFormat('HH:mm:ss:zzz')
        self.__sync_time_edit.setEnabled(False)

        self.__sync_time_button = QPushButton()
        self.__sync_time_button.setText('Sync data')
        self.__sync_time_button.setEnabled(False)
        self.__sync_time_button.clicked.connect(self.__sync_data)

        # Create the layout for the file controls
        dir_layout = QHBoxLayout()
        dir_layout.setContentsMargins(0, 0, 0, 0)
        dir_layout.addWidget(self.__dir_picker_button)
        dir_layout.addStretch(1)
        dir_layout.addWidget(self.__sync_time_label)
        dir_layout.addWidget(self.__sync_time_edit)
        dir_layout.addWidget(self.__sync_time_button)

        # Create the axis and their viewbox
        self.__x_axis_item = AxisItem('left')
        self.__y_axis_item = AxisItem('left')
        self.__z_axis_item = AxisItem('left')

        self.__x_axis_viewbox = ViewBox()
        self.__y_axis_viewbox = ViewBox()
        self.__z_axis_viewbox = ViewBox()

        # Create the widget which will display the data
        self.__graphic_view = GraphicsView(background="#ecf0f1")
        self.__graphic_layout = GraphicsLayout()
        self.__graphic_view.setCentralWidget(self.__graphic_layout)

        # Add the axis to the widget
        self.__graphic_layout.addItem(self.__x_axis_item, row=2, col=3, rowspan=1, colspan=1)
        self.__graphic_layout.addItem(self.__y_axis_item, row=2, col=2, rowspan=1, colspan=1)
        self.__graphic_layout.addItem(self.__z_axis_item, row=2, col=1, rowspan=1, colspan=1)

        self.__plot_item = PlotItem()
        self.__plot_item_viewbox = self.__plot_item.vb
        self.__graphic_layout.addItem(self.__plot_item, row=2, col=4, rowspan=1, colspan=1)

        self.__graphic_layout.scene().addItem(self.__x_axis_viewbox)
        self.__graphic_layout.scene().addItem(self.__y_axis_viewbox)
        self.__graphic_layout.scene().addItem(self.__z_axis_viewbox)

        self.__x_axis_item.linkToView(self.__x_axis_viewbox)
        self.__y_axis_item.linkToView(self.__y_axis_viewbox)
        self.__z_axis_item.linkToView(self.__z_axis_viewbox)

        self.__x_axis_viewbox.setXLink(self.__plot_item_viewbox)
        self.__y_axis_viewbox.setXLink(self.__plot_item_viewbox)
        self.__z_axis_viewbox.setXLink(self.__plot_item_viewbox)

        self.__plot_item_viewbox.sigResized.connect(self.__update_views)
        self.__x_axis_viewbox.enableAutoRange(axis=ViewBox.XAxis, enable=True)
        self.__y_axis_viewbox.enableAutoRange(axis=ViewBox.XAxis, enable=True)
        self.__z_axis_viewbox.enableAutoRange(axis=ViewBox.XAxis, enable=True)

        # Create the final layout
        self.__v_box = QVBoxLayout()
        self.__v_box.addLayout(dir_layout)
        self.__v_box.addWidget(self.__graphic_view)

        self.setLayout(self.__v_box)

        self.__restore_state()

    def __open_dir_picker(self):
        self.__shared_data.data_file_path = QFileDialog.getOpenFileUrl(self, 'Open the Hexoskin data directory',
                                                                       QDir.homePath())[0]
        if self.__shared_data.data_file_path is not None:
            try:
                self.__load_data()
                self.__add_selector_acc_gyr()
                self.__show_data('ACC_X', 'ACC_Y', 'ACC_Z')
            except FileNotFoundError:
                pass
            except UnicodeDecodeError:
                pass

    def __load_data(self):
        if self.__shared_data.data_file_path is not None:
            self.__shared_data.import_parameter()

    def __show_data(self, field1, field2, field3):
        if self.__shared_data.parameter is not None:
            # Generate the timecodes if needed
            if len(self.__shared_data.parameter['TIMECODE']) == 0:
                if self.__shared_data.sampling_rate is None:
                    result = False
                    while not result and result == 0:
                        result = self.__show_sampling_rate_picker()

                self.__shared_data.add_timecode()

            self.__x_axis_viewbox.clear()
            self.__y_axis_viewbox.clear()
            self.__z_axis_viewbox.clear()

            # Show the 3 selected fields
            self.__x_axis_viewbox.addItem(PlotCurveItem(list(map(int, self.__shared_data.parameter.get(field1))),
                                                        pen='#34495e'))
            self.__y_axis_viewbox.addItem(PlotCurveItem(list(map(int, self.__shared_data.parameter.get(field2))),
                                                        pen='#9b59b6'))
            self.__z_axis_viewbox.addItem(PlotCurveItem(list(map(int, self.__shared_data.parameter.get(field3))),
                                                        pen='#3498db'))

            self.__x_axis_item.setLabel(field1, color="#34495e")
            self.__y_axis_item.setLabel(field2, color="#9b59b6")
            self.__z_axis_item.setLabel(field3, color="#3498db")

            # Add the middle line and the bottom timecodes
            timecodes = self.__shared_data.parameter['TIMECODE']
            middle = [0] * len(timecodes)
            self.__plot_item_viewbox.addItem(PlotCurveItem(middle, pen='#000000'))
            self.__plot_item.getAxis('bottom').setTicks(
                self.__generate_time_ticks(timecodes, self.__shared_data.sampling_rate))

            # Enable the controls
            self.__sync_time_edit.setEnabled(True)
            self.__sync_time_button.setEnabled(True)
            self.__dir_picker_button.setEnabled(False)
        self.__update_views()

    def __update_views(self):
        self.__x_axis_viewbox.setGeometry(self.__plot_item_viewbox.sceneBoundingRect())
        self.__y_axis_viewbox.setGeometry(self.__plot_item_viewbox.sceneBoundingRect())
        self.__z_axis_viewbox.setGeometry(self.__plot_item_viewbox.sceneBoundingRect())

    def __generate_time_ticks(self, timecodes, rate):
        ticks = list()

        steps = [rate * 30, rate * 15, rate * 5, rate]
        for step in steps:
            temp = list()
            i = step
            while i in range(len(timecodes)):
                temp.append((i, timecodes[i].strftime('%H:%M:%S:') + str(int(timecodes[i].microsecond / 1000))))
                i += step
            ticks.append(temp)

        return ticks

    def __sync_data(self):
        self.__shared_data.data_sync = self.__sync_time_edit.text()
        self.__shared_data.update_sync.emit()

    def __show_sampling_rate_picker(self) -> bool:
        self.__shared_data.sampling_rate, result = QInputDialog.getInt(self, 'Set sampling rate value', 'Sampling rate')
        return result

    def __add_selector_acc_gyr(self):
        if 'GYR_X' in self.__shared_data.parameter.keys() or 'GYR_Y' in self.__shared_data.parameter.keys() \
                or 'GYR_Z' in self.__shared_data.parameter.keys():
            show_acc = QPushButton()
            show_acc.setText('Show Accelerometer Axis')
            show_acc.clicked.connect(self.__show_acc)

            show_gyr = QPushButton()
            show_gyr.setText('Show Gyroscope Axis')
            show_gyr.clicked.connect(self.__show_gyr)

            layout = QHBoxLayout()
            layout.addWidget(show_acc)
            layout.addWidget(show_gyr)
            layout.addStretch(1)

            self.__v_box.addLayout(layout)

    def __show_acc(self):
        self.__show_data('ACC_X', 'ACC_Y', 'ACC_Z')

    def __show_gyr(self):
        self.__show_data('GYR_X', 'GYR_Y', 'GYR_Z')

    def __restore_state(self):
        if self.__shared_data.parameter is not None:
            self.__add_selector_acc_gyr()
            self.__show_data('ACC_X', 'ACC_Y', 'ACC_Z')
            print('trigger reimport')
        if self.__shared_data.data_sync is not None:
            text_time = self.__shared_data.data_sync.split(':')
            time = QTime()
            time.setHMS(int(text_time[0]), int(text_time[1]), int(text_time[2]), int(text_time[3]))
            self.__sync_time_edit.setTime(time)
예제 #15
0
class PlotsWidget(QtGui.QWidget):
    """
    """
    def __init__(self):
        QtGui.QWidget.__init__(self)

        self.data = {}

        self.gv = pg.GraphicsView()
        self.gl = GraphicsLayout()

        self.gv.setCentralItem(self.gl)

        self.layout = QtGui.QGridLayout()
        self.layout.addWidget(self.gv, 1, 1, 1, 1)

        self.setLayout(self.layout)

    def clear(self):
        for _, pl in np.ndenumerate(self.pl):
            pl.clear()

    def add_data(self, data, label='plot'):
        self.data[label] = {'x': data['time'], 'y': data['data']}

    def setup_plots(self, nrow, ncol, title=None, clickable=True):
        self.nrow = nrow
        self.ncol = ncol

        if title is not None:
            title_k = list(title.keys())

        self.pl = np.zeros((nrow, ncol), dtype=object)
        for i in range(nrow):
            for j in range(ncol):
                plot_title = "{}: {}".format(title_k[i], title[title_k[i]][j]) if title is not None else None
                self.pl[i, j] = self.gl.addPlot(title=plot_title, row=i, col=j, clickable=clickable)

    def plot(self, row_id, col_id, label=None, offset=(350, 30), **kwargs):
        self.pl[row_id, col_id].plot(**kwargs)

    def bar(self, row_id, col_id, label=None, offset=(350, 30), **kwargs):
        """
        Plotting the bar graphs onto Plot Items.
        """
        bar = BarGraph(label=label, **kwargs)
        self.pl[row_id, col_id].addItem(bar)

        bar.barClicked.connect(self.clickedBar)

    def imshow(self, row_id, col_id, im_data, label=None, set_pos=None, scale=None, set_levels=None,
               set_aspect_locked=False, invert_y=False, invert_x=False):

        self.lut = getLUT()
        img = pg.ImageItem()
        self.pl[row_id, col_id].addItem(img)

        # Update ImageView with new data.
        img_item = self.pl[row_id, col_id].items[-1]
        img_item.setImage(im_data, lut=self.lut)

        if set_pos is not None:
            img_item.setPos(*set_pos)

        if scale is not None:
            img_item.scale(*scale)

        if set_levels is not None:
            img_item.setLevels(set_levels)

        self.pl[row_id, col_id].setAspectLocked(set_aspect_locked)

        self.pl[row_id, col_id].invertY(invert_y)
        self.pl[row_id, col_id].invertX(invert_x)

    def clickedBar(self, message):

        tvec = self.data['plot']['x'][message._last_index]
        data = self.data['plot']['y'][message._last_index]

        dlg = QtWidgets.QDialog()
        dlg.setMinimumSize(800, 600)
        dlg.setLayout(QtWidgets.QVBoxLayout(dlg))
        glw = pg.GraphicsLayoutWidget(parent=dlg)
        dlg.layout().addWidget(glw)
        y_range = (np.min(data), np.max(data))
        x_range = (np.min(tvec), np.min(tvec) + 4)
        for ch in range(data.shape[0]):
            plt = glw.addPlot(row=ch, col=0)
            pen = QtGui.QColor(THEMES['dark']['pencolors'][ch])
            curve = plt.plot(x=tvec, y=data[ch, :], pen=pen)

            # curve = plt.plot(x=tvec, y= data[ch, :], name=self.data['labels'][ch], pen=pen)
            plt.setYRange(*y_range)
            plt.setXRange(*x_range)

        dlg.exec_()