1# -*- coding: utf-8 -*- 2 3import numpy as np 4from PySide2.QtCore import QAbstractTableModel, QModelIndex, Qt 5from PySide2.QtGui import QColor 6 7 8class Model(QAbstractTableModel): 9 def __init__(self): 10 QAbstractTableModel.__init__(self) 11 12 self.ncols = 4 13 nrows = 1 # initial number of rows 14 15 self.data_x = np.zeros(nrows) 16 self.data_xerr = np.zeros(nrows) 17 self.data_y = np.zeros(nrows) 18 self.data_yerr = np.zeros(nrows) 19 20 def rowCount(self, parent=QModelIndex()): 21 return self.data_x.size 22 23 def columnCount(self, parent=QModelIndex()): 24 return self.ncols 25 26 def flags(self, index): 27 return Qt.ItemIsEditable | Qt.ItemIsEnabled | Qt.ItemIsSelectable 28 29 def setData(self, index, value, role): 30 if index.isValid() and role == Qt.EditRole: 31 column = index.column() 32 row = index.row() 33 34 try: 35 float_value = float(value) 36 37 if column == 0: 38 self.data_x[row] = float_value 39 40 if column == 1: 41 self.data_xerr[row] = float_value 42 43 if column == 2: 44 self.data_y[row] = float_value 45 46 if column == 3: 47 self.data_yerr[row] = float_value 48 49 self.dataChanged.emit(index, index) 50 51 return True 52 except ValueError: 53 return False 54 55 return False 56 57 def headerData(self, section, orientation, role): 58 if role != Qt.DisplayRole: 59 return None 60 61 if orientation == Qt.Horizontal: 62 return ("x", "xerr", "y", "yerr")[section] 63 64 return "{}".format(section) 65 66 def data(self, index, role): 67 if role == Qt.BackgroundRole: 68 return QColor(Qt.transparent) 69 70 elif role == Qt.TextAlignmentRole: 71 return Qt.AlignRight 72 73 elif role == Qt.DisplayRole: 74 column = index.column() 75 row = index.row() 76 77 if column == 0: 78 return "{0:.6e}".format(self.data_x[row]) 79 80 if column == 1: 81 return "{0:.6e}".format(self.data_xerr[row]) 82 83 if column == 2: 84 return "{0:.6e}".format(self.data_y[row]) 85 86 if column == 3: 87 return "{0:.6e}".format(self.data_yerr[row]) 88 89 def sort(self, column, order): 90 if column == 0: 91 sorted_idx = self.data_x.argsort() 92 93 if order == Qt.SortOrder.DescendingOrder: 94 sorted_idx = sorted_idx[::-1] 95 96 self.beginResetModel() 97 98 self.data_x = self.data_x[sorted_idx] 99 self.data_xerr = self.data_xerr[sorted_idx] 100 self.data_y = self.data_y[sorted_idx] 101 self.data_yerr = self.data_yerr[sorted_idx] 102 103 self.endResetModel() 104 105 first_index = self.index(0, 0) 106 last_index = self.index(self.rowCount() - 1, self.columnCount() - 1) 107 108 self.dataChanged.emit(first_index, last_index) 109 110 def append_row(self): 111 self.beginInsertRows(QModelIndex(), self.data_x.size, self.data_x.size) 112 113 self.data_x = np.append(self.data_x, 0) 114 self.data_xerr = np.append(self.data_xerr, 0) 115 self.data_y = np.append(self.data_y, 0) 116 self.data_yerr = np.append(self.data_yerr, 0) 117 118 self.endInsertRows() 119 120 first_index = self.index(0, 0) 121 last_index = self.index(self.rowCount() - 1, self.columnCount() - 1) 122 123 self.dataChanged.emit(first_index, last_index) 124 125 def remove_rows(self, index_list): 126 index_list.sort(reverse=True) 127 128 self.data_x = np.delete(self.data_x, index_list) 129 self.data_xerr = np.delete(self.data_xerr, index_list) 130 self.data_y = np.delete(self.data_y, index_list) 131 self.data_yerr = np.delete(self.data_yerr, index_list) 132 133 for index in index_list: 134 self.beginRemoveRows(QModelIndex(), index, index) 135 self.endRemoveRows() 136 137 first_index = self.index(index_list[-1], 0) 138 last_index = self.index(index_list[0], 3) 139 140 self.dataChanged.emit(first_index, last_index) 141