dylab/HelperClasses/Plotting/WidgetQuickWaterFlowPlot.py

class ArrayDataCombo(ExtendedCombo):

    def __init__(self, h5_paths, **kwargs):

        super().__init__(**kwargs)
        self.update_model(h5_paths)

    def update_model(self, h5_paths):
        results_array_labels = set([])
        for h5_path in h5_paths:
            with h5py.File(h5_path, 'r') as h5_file:
                analysis_names = h5_file['results'].keys()

                for analysis_name in analysis_names:
                    for key in h5_file['results'][analysis_name].keys():
                        results_array_labels.add(analysis_name + ',' + key)

        model = QtGui.QStandardItemModel()

        item = QtGui.QStandardItem('shot number')
        model.setItem(0, 0, item)

        for i, idx in enumerate(results_array_labels):
            item = QtGui.QStandardItem(idx)
            model.setItem(i, 0, item)

        self.setModel(model)

    def get_idx(self):
        return tuple(str(self.currentText()).split(','))


class QuickWaterfallPlot(QuickDataPlot):

    def __init__(self, *args, **kwargs):

        super().__init__(*args, **kwargs)

        self.img = pg.ImageItem()
        self.plot.addItem(self.img)

        # Isocurve draplotsg
        self.iso = pg.IsocurveItem(level=1000, pen=color_palette[2])
        self.iso.setParentItem(self.img)
        self.iso.setZValue(5)

        # Contrast/color control
        self.hist = pg.HistogramLUTItem()
        self.hist.setImageItem(self.img)
        self.plots.addItem(self.hist)

        # Draggable line for setting isocurve level
        self.isoLine = pg.InfiniteLine(angle=0, movable=True, pen=color_palette[2])
        self.hist.vb.addItem(self.isoLine)
        self.hist.vb.setMouseEnabled(y=False)  # makes user interaction a little easier
        self.isoLine.setValue(1000)
        self.isoLine.setZValue(1000)  # bring iso line above contrast controls
        self.isoLine.sigDragged.connect(self.updateIsocurve)

        # Monkey-patch the image to use our custom hover function.
        # This is generally discouraged (you should subclass ImageItem instead),
        # but it works for a very simple use like this.
        self.img.hoverEvent = self.imageHoverEvent

        self.img.translate(-0.5, -0.5)

        self.scalex = 1
        self.scaley = 1

        self.cx = 0
        self.cy = 0

        self.nplots = 0

        self.table.setColumnCount(3)
        self.table.setRowCount(2)
        self.table.setColumnWidth(0, 200)
        self.table.setColumnWidth(1, 150)
        self.table.setColumnWidth(2, 150)
        self.table.setHorizontalHeaderLabels(['xvalue', 'yarray', 'zarray'])

        self.combox = NumericDataCombo(self.ap.df)
        self.comboy = ArrayDataCombo(self.ap.h5_paths)
        self.comboz = ArrayDataCombo(self.ap.h5_paths)

        self.table.setCellWidget(0, 0, self.combox)
        self.table.setCellWidget(0, 1, self.comboy)
        self.table.setCellWidget(0, 2, self.comboz)

        self.mk_buttons()

    def mk_buttons(self):
        self.bt_update = QPushButton('Update', self)
        self.bt_update.clicked.connect(self.update_from_h5)
        self.table.setCellWidget(1, 1, self.bt_update)

    def update_data_extractor(self):

        idxx = self.combox.get_idx()
        idxy = self.comboy.get_idx()
        idxz = self.comboz.get_idx()

        if idxx not in self.data_extractor.data_extractors:
            if idxx[0] == 'shot number':
                self.data_extractor[idxx] = EmptyDataExtractor()
            else:
                self.data_extractor[idxx] = SingleDataExtractor(idxx)

        if idxy not in self.data_extractor.data_extractors:
            self.data_extractor[idxy] = ArrayDataExtractor(idxy)

        if idxz not in self.data_extractor.data_extractors:
            self.data_extractor[idxz] = ArrayDataExtractor(idxz)

        self.data_extractor.clean_children([idxx, idxy, idxz])
        self.data_extractor.clean_memory(self.ap.h5_paths)

    def update(self, data=None):
        xs = np.array([])
        ys = np.array([])
        zs = np.array([])

        idxx = self.combox.get_idx()
        idxy = self.comboy.get_idx()
        idxz = self.comboz.get_idx()

        for i, h5_path in enumerate(self.ap.h5_paths_selected):

            data = self.data_extractor.get_data(h5_path)[0]

            if idxx[0] == 'shot number':
                x = self.ap.shotselector.get_selected_indices()[0][i]
            else:
                x = data[idxx]

            y = data[idxy]
            z = data[idxz]

            if y is not None and z is not None:
                xs = np.append(xs, x * np.ones_like(y))
                ys = np.append(ys, y)
                zs = np.append(zs, z)
            elif y is not None:
                xs = np.append(xs, x * np.ones_like(y))
                ys = np.append(ys, y)
                zs = np.append(zs, y * np.nan)
            elif z is not None:
                xs = np.append(xs, x * np.ones_like(z))
                ys = np.append(ys, z * np.nan)
                zs = np.append(zs, z)

        # here we don't want to assume that the data is on a grid
        # this can happen if the parameters where changed between two sweeps
        xi = np.linspace(xs.min(), xs.max(), 200)
        yi = np.linspace(ys.min(), ys.max(), 200)

        Xi, Yi = np.meshgrid(xi, yi)

        Zi = griddata((xs, ys), zs, (Xi, Yi), 'nearest')

        self.img.setImage(Zi.T)
        self.iso.setData(Zi.T)
        self.data_img = Zi

        # set position and scale of image
        newscalex = xi[1] - xi[0]
        newscaley = yi[1] - yi[0]

        transx = (xi[0] - (self.cx - 0.5 * self.scalex)) / newscalex - 0.5
        transy = (yi[0] - (self.cy - 0.5 * self.scaley)) / newscaley - 0.5

        self.img.scale(newscalex / self.scalex, newscaley / self.scaley)
        self.img.translate(transx, transy)

        self.scalex = newscalex
        self.scaley = newscaley

        self.cx = xi[0]
        self.cy = yi[0]

    def updateIsocurve(self):
        self.iso.setLevel(self.isoLine.value())

    def imageHoverEvent(self, event):
        """Show the position, pixel, and value under the mouse cursor.
        """
        if event.isExit():
            self.plot.setTitle("")
            return
        pos = event.pos()
        i, j = pos.y(), pos.x()
        i = int(np.clip(i, 0, self.data_img.shape[0] - 1))
        j = int(np.clip(j, 0, self.data_img.shape[1] - 1))
        val = self.data_img[i, j]
        ppos = self.img.mapToParent(pos)
        x, y = ppos.x(), ppos.y()
        self.plot.setTitle("pos: (%0.1f, %0.1f) value: %g" % (x, y, val))