# 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,
help='Name of the currently selected array.')
arf = plot.ArrayResponseFunctionPlotState.T(
default=plot.ArrayResponseFunctionPlotState.D(),
help='Array response function plot settings.')
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,
plot_args=(self.state.arf,))
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.arf.plot.set_array(array, info)