import sys import os import logging import csv from PyQt5 import QtCore, QtGui import numpy as np logger = logging.getLogger("helical") from PyQt5.QtWidgets import ( QTableWidget, QApplication, QMainWindow, QTableWidgetItem, QFileDialog, QHeaderView, QAbstractItemView, QMenu) from PyQt5.QtCore import Qt, QFileInfo, pyqtSignal, QPoint, pyqtSlot DATA_ITEM = 1 START_OMEGA_0 = 0 START_OMEGA_120 = 1 START_OMEGA_240 = 2 STOP_OMEGA_0 = 3 STOP_OMEGA_120 = 4 STOP_OMEGA_240 = 5 ROLE_XTAL_START = 1 + Qt.UserRole ROLE_XTAL_END = 2 + Qt.UserRole class HelicalTableWidget(QTableWidget): gonioMoveRequest = pyqtSignal(float, float, float, float, float) def __init__(self): super().__init__() self.check_change = True self._scanAngularStep = 0.1 self._scanHorizontalCount = 5 self._scanVerticalCount = 10 self._current_xtal = None self._start_angle = 0.0 self.init_ui() def startAngle(self) -> float: return self._start_angle @pyqtSlot(float) def setStartAngle(self, angle:float): self._start_angle = angle def init_ui(self): self.setColumnCount(6) self.setEditTriggers(QAbstractItemView.NoEditTriggers) self.resizeColumnsToContents() self.setSelectionMode(QAbstractItemView.SingleSelection) self.setSelectionBehavior(QAbstractItemView.SelectItems) labels = [ "0\nbx, bz\n(mm)", "120\nbx, bz\n(mm)", "240\nbx, bz\n(mm)", "0\nbx, bz\n(mm)", "120\nbx, bz\n(mm)", "240\nbx, bz\n(mm)", ] self.setHorizontalHeaderLabels(labels) self.verticalHeader().resizeSections(QHeaderView.ResizeToContents) self.horizontalHeader().setSectionResizeMode(QHeaderView.Stretch) self.cellChanged.connect(self.c_current) self.show() def display_coords(self, x, y): row = self.rowAt(y) col = self.columnAt(x) coords = self.gonio_coords_at(row, col) omega = col % 3 * 120 data = coords[omega] logger.info("gonio at {}, {}: {}".format(row, col, data)) def goto_gonio_position(self, x, y): row = self.rowAt(y) col = self.columnAt(x) coords = self.gonio_coords_at(row, col) omega = col % 3 * 120 data = coords[omega] logger.info("move gonio to: {}".format(row, col, data)) self.gonioMoveRequest.emit(*data) def remove_xtal(self, x, y): row = self.rowAt(y) col = self.columnAt(x) logger.info("rowCount() = {}, removing row: {}".format(self.rowCount(), row)) self.removeRow(row) if self.rowCount() == 0: self._current_xtal = None logger.info("rowCount() = {}".format(self.rowCount())) def contextMenuEvent(self, event): logger.info(event.pos()) x = event.pos().x() y = event.pos().y() menu = QMenu(self) gotoAction = menu.addAction("Move Gonio Here") gotoAction.triggered.connect(lambda b: self.goto_gonio_position(x, y)) removeAction = menu.addAction("Remove this line") removeAction.triggered.connect(lambda b: self.remove_xtal(x, y)) menu.popup(QtGui.QCursor.pos()) def setScanHorizontalCount(self, count: int): logger.debug("horizontal count: {}".format(count)) self._scanHorizontalCount = count def scanHorizontalCount(self) -> int: return self._scanHorizontalCount def setScanVerticalCount(self, count: int): logger.debug("vertical count: {}".format(count)) self._scanVerticalCount = count def scanVerticalCount(self) -> int: return self._scanVerticalCount def setScanAngularStep(self, step: float): self._scanAngularStep = step def scanAngularStep(self) -> float: return self._scanAngularStep def scanTotalRange(self) -> float: return self._scanVerticalCount * self._scanHorizontalCount * self._scanAngularStep def gonio_coords_at(self, row, col): role = ROLE_XTAL_START if col < 3 else ROLE_XTAL_END coords = self.item(row, DATA_ITEM).data(role) return coords def add_xtal(self): row = self.rowCount() self.setRowCount(row + 1) self._current_xtal = row print(row) for n in range(self.columnCount()): self.setItem(row, n, QTableWidgetItem("unset")) def set_xtal_start(self, data: list): """sets data to start position data = [fx, fy, bx, bz, omega] """ self.set_data_point(ROLE_XTAL_START, data) def set_xtal_end(self, data: list): self.set_data_point(ROLE_XTAL_END, data) def set_data_point(self, role: int, data: list): if self._current_xtal is None: self.add_xtal() row, col = self._current_xtal, DATA_ITEM try: coords = self.item(row, DATA_ITEM).data(role) except: print("empty coords; initializing...") coords = None if coords is None: coords = {} fx, fy, bx, bz, omega = data o = int(omega) coords[o] = data print("coords = {}".format(coords)) if role == ROLE_XTAL_END: col = 3 else: col = 0 for n in range(3): omega = n * 120 try: fx, fy, bx, bz, omega = coords[omega] info = "{bx:.3f} mm\n{bz:.3f} mm".format(bx=bx, bz=bz) except: info = "unset" # self.item(row, col + n).setData(Qt.DisplayRole, info) self.item(row, col + n).setText(info) self.item(row, col + n).setToolTip(info) self.item(row, DATA_ITEM).setData(role, coords) def c_current(self): if self.check_change: row = self.currentRow() col = self.currentColumn() value = self.item(row, col) # value = value.text() def load_datapoints(self): self.check_change = False path = QFileDialog.getOpenFileName( self, "Open CSV", os.getenv("HOME"), "CSV(*.csv)" ) if path[0] != "": with open(path[0], newline="") as csv_file: self.setRowCount(0) self.setColumnCount(10) my_file = csv.reader(csv_file, delimiter=",", quotechar="|") for row_data in my_file: row = self.rowCount() self.insertRow(row) if len(row_data) > 10: self.setColumnCount(len(row_data)) for column, stuff in enumerate(row_data): item = QTableWidgetItem(stuff) self.setItem(row, column, item) self.check_change = True def get_data(self, as_numpy=False): """return a list of tuples with all defined data [ (is_fiducial(boolean), x, y, gonio_x, gonio_y),... ] """ data = [] for row in range(self.rowCount()): start = self.item(row, DATA_ITEM).data(ROLE_XTAL_START) end = self.item(row, DATA_ITEM).data(ROLE_XTAL_END) data.append([start, end]) # if as_numpy: # data = np.asarray(data) return data if __name__ == "__main__": class Sheet(QMainWindow): def __init__(self): super().__init__() self.form_widget = HelicalTableWidget() self.setCentralWidget(self.form_widget) self.show() app = QApplication(sys.argv) sheet = Sheet() sys.exit(app.exec_())