Source code for pyrocko.gato.gui.browser

# https://pyrocko.org - GPLv3
#
# The Pyrocko Developers, 21st Century
# ---|P------/S----------~Lg----------


from pyrocko.guts import String
from pyrocko.gui.qt_compat import qw, qc
from pyrocko.gui.util import SmartplotFrame
from pyrocko.gui import talkie
from pyrocko.gui.state import state_bind
from pyrocko import gato
from pyrocko.gato import plot


class ArrayInventory(qc.QAbstractTableModel):
    def __init__(self):
        qc.QAbstractTableModel.__init__(self)
        self.columns = [
            'name', 'type', 'comment', 'n_codes_nsl', 'n_codes', 'str_tmin',
            'str_tmax', 'str_codes_nsl_by_channels']
        self.column_titles = [
            'Name', 'Type', 'Comment', '# Sites', '# Channels', 'Start Date',
            'End Date', 'Channel Group: # Sites']
        self.column_sizes = [100] * len(self.columns)
        self.column_sizes[2] = 350

        self.arrays = []
        self.array_infos = []

    def add_arrays(self, arrays):
        self.beginInsertRows(
             qc.QModelIndex(),
             len(self.arrays),
             len(self.arrays) + len(arrays) - 1)

        self.arrays.extend(arrays)
        self.array_infos.extend([None] * len(arrays))
        self.endInsertRows()
        self.update_array_infos()

    def get_array(self, index):
        return self.arrays[index.row()]

    def get_array_and_info_by_name(self, name):
        for array, info in zip(self.arrays, self.array_infos):
            if array.name == name:
                return array, info

        return None

    def _get_index_by_name(self, name):
        for i, array in enumerate(self.arrays):
            if name == array.name:
                return i

        return None

    def get_index_by_name(self, name):
        i = self._get_index_by_name(name)
        if i is None:
            return qc.QModelIndex()

        return self.index(i, 0)

    def get_array_info(self, index):
        return self.array_infos[index.row()]

    def update_array_infos(self):
        from .main import get_squirrel
        sq = get_squirrel()
        if sq is None:
            return

        for i in range(len(self.arrays)):
            if self.array_infos[i] is None:
                self.array_infos[i] = self.arrays[i].get_info(sq)
                istart = self.index(3, i)
                istop = self.index(len(self.columns)-1, i)
                self.dataChanged.emit(istart, istop)

    def rowCount(self, parent):
        return len(self.arrays)

    def columnCount(self, parent):
        return len(self.columns)

    def headerData(self, section, orientation, role):
        if orientation == qc.Qt.Horizontal:
            if role == qc.Qt.DisplayRole:
                return qc.QVariant(self.column_titles[section])
            elif role == qc.Qt.SizeHintRole:
                return qc.QSize(10, 20)

        elif orientation == qc.Qt.Vertical:
            if role == qc.Qt.DisplayRole:
                return qc.QVariant(str(section))

        return qc.QVariant()

    def data(self, index, role):
        irow = index.row()
        icol = index.column()
        column = self.columns[icol]
        array = self.arrays[irow]
        array_info = self.array_infos[irow]

        obj = array if icol < 3 else array_info
        if obj is None:
            return qc.QVariant()

        if role == qc.Qt.BackgroundRole:
            return qc.QVariant()

        elif role == qc.Qt.ForegroundRole:
            return qc.QVariant()

        elif role in (qc.Qt.DisplayRole, qc.Qt.UserRole):
            return qc.QVariant(getattr(obj, column))

        else:
            return qc.QVariant()


[docs]class ArrayBrowserState(talkie.TalkieRoot): current_array_name = String.T(optional=True)
class ArrayBrowser(qw.QFrame, talkie.TalkieConnectionOwner): def __init__(self, *args, **kwargs): qw.QFrame.__init__(self, *args, **kwargs) talkie.TalkieConnectionOwner.__init__(self) self.state = ArrayBrowserState() layout = qw.QGridLayout() self.setLayout(layout) self.inventory = inventory = ArrayInventory() arrays = gato.get_named_arrays().values() inventory.add_arrays(arrays) array_table = qw.QTableView() array_table.setModel(inventory) array_table.setShowGrid(False) array_table.setSelectionBehavior(qw.QAbstractItemView.SelectRows) array_table.verticalHeader().hide() selection = qc.QItemSelectionModel(inventory) array_table.setSelectionModel(selection) def update_state(widget, state): current_array = self.inventory.get_array(selection.currentIndex()) self.state.current_array_name \ = current_array.name if current_array else None def update_widget(state, widget): index = self.inventory.get_index_by_name( self.state.current_array_name) selection.setCurrentIndex(index, qc.QItemSelectionModel.Current) state_bind( self, self.state, ['current_array_name'], update_state, array_table, [selection.currentRowChanged], update_widget) header = array_table.horizontalHeader() for i_s, s in enumerate(inventory.column_sizes): header.resizeSection(i_s, s) header.setStretchLastSection(True) self.array_table = array_table ny = 0 layout.addWidget(array_table, ny, 0, 1, 2) ny += 1 self._geometry_items = [] geometry = SmartplotFrame(plot_cls=plot.GeometryPlot) layout.addWidget(geometry, ny, 0, 1, 1) geometry.draw() self.geometry = geometry arf = SmartplotFrame(plot_cls=plot.ArrayResponseFunctionPlot) layout.addWidget(arf, ny, 1, 1, 1) self.arf = arf self.talkie_connect( self.state, ['current_array_name'], self.update_current_array) def update_array_infos(self): self.inventory.update_array_infos() def update_current_array(self, *args): self.current_array, self.current_array_info = \ self.inventory.get_array_and_info_by_name( self.state.current_array_name) self.update_plots() def update_plots(self): array = self.current_array info = self.current_array_info self.geometry.plot.set_array(array, info) self.geometry.draw() self.arf.plot.set_array(array, info) self.arf.draw()