From 25ffd27e5535513370a8431f71a6dcfb13feae77 Mon Sep 17 00:00:00 2001 From: JoschD <26184899+JoschD@users.noreply.github.com> Date: Tue, 14 Nov 2023 21:16:40 +0100 Subject: [PATCH] pretend to load measurements, parse info, color and tooltips --- omc3_gui/segment_by_segment/controller.py | 15 +- .../segment_by_segment/measurement_model.py | 165 ++++++++++++++++++ .../segment_by_segment/measurement_view.py | 0 omc3_gui/segment_by_segment/model.py | 78 +++++---- omc3_gui/segment_by_segment/view.py | 109 ++++++++---- .../utils/{dialogs.py => file_dialogs.py} | 9 + 6 files changed, 303 insertions(+), 73 deletions(-) create mode 100644 omc3_gui/segment_by_segment/measurement_model.py create mode 100644 omc3_gui/segment_by_segment/measurement_view.py rename omc3_gui/utils/{dialogs.py => file_dialogs.py} (82%) diff --git a/omc3_gui/segment_by_segment/controller.py b/omc3_gui/segment_by_segment/controller.py index 4fc5336..ae93ee0 100644 --- a/omc3_gui/segment_by_segment/controller.py +++ b/omc3_gui/segment_by_segment/controller.py @@ -1,10 +1,12 @@ from pathlib import Path +from typing import List from omc3_gui.utils.base_classes import Controller -from omc3_gui.utils.dialogs import OpenDirectoriesDialog, OpenDirectoryDialog +from omc3_gui.utils.file_dialogs import OpenDirectoriesDialog, OpenDirectoryDialog from omc3_gui.segment_by_segment.view import SbSWindow -from omc3_gui.segment_by_segment.model import Measurement, Settings +from omc3_gui.segment_by_segment.model import Settings from qtpy.QtCore import Qt, Signal, Slot from qtpy.QtWidgets import QFileDialog +from omc3_gui.segment_by_segment.measurement_model import OpticsMeasurement import logging LOG = logging.getLogger(__name__) @@ -21,7 +23,7 @@ def __init__(self): self._last_selected_optics_path = None - def add_measurement(self, measurement: Measurement): + def add_measurement(self, measurement: OpticsMeasurement): self._view.get_measurement_list().add_item(measurement) @@ -44,7 +46,10 @@ def open_measurements(self): for filename in filenames: self._last_selected_optics_path = filename.parent LOG.debug(f"User selected: {filename}") - loaded_measurements.add_item(Measurement(filename)) + optics_measurement = OpticsMeasurement.from_path(filename) + try: + loaded_measurements.add_item(optics_measurement) + except ValueError as e: + LOG.error(str(e)) - diff --git a/omc3_gui/segment_by_segment/measurement_model.py b/omc3_gui/segment_by_segment/measurement_model.py new file mode 100644 index 0000000..2627de9 --- /dev/null +++ b/omc3_gui/segment_by_segment/measurement_model.py @@ -0,0 +1,165 @@ + +from dataclasses import dataclass, fields +from pathlib import Path +import types +from typing import Any, Dict, List, Sequence, Union, Optional, ClassVar +from accwidgets.graph import StaticPlotWidget +import pyqtgraph as pg +from omc3.segment_by_segment.segments import Segment +from omc3.definitions.optics import OpticsMeasurement as OpticsMeasurementCollection +from omc3.optics_measurements.constants import MODEL_DIRECTORY, KICK_NAME, PHASE_NAME, BETA_NAME, EXT +from omc3.model.constants import TWISS_DAT +from tfs.reader import read_headers +import logging + + +SEQUENCE = "SEQUENCE" +DATE = "DATE" +LHC_MODEL_YEARS = (2012, 2015, 2016, 2017, 2018, 2022, 2023) # TODO: get from omc3 + +FILES_TO_LOOK_FOR = (f"{name}{plane}" for name in (KICK_NAME, PHASE_NAME, BETA_NAME) for plane in ("x", "y")) + +LOGGER = logging.getLogger(__name__) + +@dataclass +class OpticsMeasurement: + measurement_dir: Path # Path to the optics-measurement folder + model_dir: Path = None # Path to the model folder + accel: str = None # Name of the accelerator + output_dir: Optional[Path] = None # Path to the sbs-output folder + elements: Optional[Dict[str, Segment]] = None # List of elements + segments: Optional[Dict[str, Segment]] = None # List of segments + year: Optional[str] = None # Year of the measurement (accelerator) + ring: Optional[int] = None # Ring of the accelerator + beam: Optional[int] = None # LHC-Beam (not part of SbS input!) + + DEFAULT_OUTPUT_DIR: ClassVar[str] = "sbs" + + def __post_init__(self): + if self.output_dir is None: + self.output_dir = self.measurement_dir / self.DEFAULT_OUTPUT_DIR + + def display(self) -> str: + return str(self.measurement_dir.name) + + def tooltip(self) -> str: + parts = ( + ("Optics Measurement", self.measurement_dir), + ("Model", self.model_dir), + ("Accelerator", self.accel), + ("Beam", self.beam), + ("Year", self.year), + ("Ring", self.ring), + ) + l = max(len(name) for name, _ in parts) + return "\n".join(f"{name:{l}s}: {value}" for name, value in parts if value is not None) + + @classmethod + def from_path(cls, path: Path) -> "OpticsMeasurement": + """ Creates an OpticsMeasurement from a folder, by trying + to parse information from the data in the folder. + + Args: + path (Path): Path to the folder. + + Returns: + OpticsMeasurement: OpticsMeasurement instance. + """ + model_dir = None + info = {} + try: + model_dir = _parse_model_dir_from_optics_measurement(path) + except FileNotFoundError as e: + LOGGER.error(str(e)) + else: + info = _parse_info_from_model_dir(model_dir) + + meas = cls(measurement_dir=path, model_dir=model_dir, **info) + if ( + any(getattr(meas, name) is None for name in ("model_dir", "accel", "output_dir")) + or (meas.accel == 'lhc' and (meas.year is None or meas.beam is None)) + or (meas.accel == 'psb' and meas.ring is None) + ): + LOGGER.error(f"Info parsed from measurement folder '{path!s}' is incomplete. Adjust manually!!") + return meas + + +def _parse_model_dir_from_optics_measurement(measurement_path: Path) -> Path: + """Tries to find the model directory in the headers of one of the optics measurement files. + + Args: + measurement_path (Path): Path to the folder. + + Returns: + Path: Path to the (associated) model directory. + """ + LOGGER.debug(f"Searching for model dir in {measurement_path!s}") + for file_name in FILES_TO_LOOK_FOR: + LOGGER.debug(f"Checking {file_name!s} for model dir.") + try: + headers = read_headers((measurement_path / file_name).with_suffix(EXT)) + except FileNotFoundError: + LOGGER.debug(f"{file_name!s} not found in {measurement_path!s}.") + else: + if MODEL_DIRECTORY in headers: + LOGGER.debug(f"{MODEL_DIRECTORY!s} found in {file_name!s}!") + break + + LOGGER.debug(f"{MODEL_DIRECTORY!s} not found in {file_name!s}.") + else: + raise FileNotFoundError(f"Could not find '{MODEL_DIRECTORY}' in any of {FILES_TO_LOOK_FOR!r} in {measurement_path!r}") + path = Path(headers[MODEL_DIRECTORY]) + LOGGER.debug(f"Associated model dir found: {path!s}") + return path + + +def _parse_info_from_model_dir(model_dir: Path) -> Dict[str, Any]: + """ Checking twiss.dat for more info about the accelerator. + + Args: + model_dir (Path): Path to the model-directory. + + Returns: + Dict[str, Any]: Containing the additional info found (accel, beam, year, ring). + """ + result = {} + + try: + headers = read_headers(model_dir / TWISS_DAT) + except FileNotFoundError as e: + LOGGER.debug(str(e)) + return result + + sequence = headers.get(SEQUENCE) + if sequence is not None: + sequence = sequence.lower() + if "lhc" in sequence: + result['accel'] = "lhc" + result['beam'] = int(sequence[-1]) + result['year'] = _get_lhc_model_year(headers.get(DATE)) + elif "psb" in sequence: + result['accel'] = "psb" + result['ring'] = int(sequence[-1]) + else: + result['accel'] = sequence + LOGGER.debug(f"Associated info found in model dir '{model_dir!s}':\n {result!s}") + return result + + +def _get_lhc_model_year(date: Union[str, None]) -> Union[str, None]: + """ Parses the year from the date in the LHC twiss.dat file + and tries to find the closest model-year.""" + if date is None: + return None + try: + found_year = int(f"20{date.split('/')[-1]}") + except ValueError: + LOGGER.debug(f"Could not parse year from '{date}'!") + return None + + for year in sorted(LHC_MODEL_YEARS, reverse=True): + if year <= found_year: + return str(year) + + LOGGER.debug(f"Could not parse year from '{date}'!") + return None \ No newline at end of file diff --git a/omc3_gui/segment_by_segment/measurement_view.py b/omc3_gui/segment_by_segment/measurement_view.py new file mode 100644 index 0000000..e69de29 diff --git a/omc3_gui/segment_by_segment/model.py b/omc3_gui/segment_by_segment/model.py index 5b9b6c0..34fa9fe 100644 --- a/omc3_gui/segment_by_segment/model.py +++ b/omc3_gui/segment_by_segment/model.py @@ -1,30 +1,16 @@ -from dataclasses import dataclass -from pathlib import Path +import enum import types -from typing import Any, Dict, List, Sequence, Union -from accwidgets.graph import StaticPlotWidget +from dataclasses import dataclass, fields +from pathlib import Path +from typing import Any, ClassVar, Dict, List, Optional, Sequence, Union + import pyqtgraph as pg +from accwidgets.graph import StaticPlotWidget from omc3.segment_by_segment.segments import Segment +from qtpy import QtCore, QtWidgets +from qtpy.QtCore import QModelIndex, Qt -from typing import List -from qtpy import QtCore -from qtpy.QtCore import Qt - - -@dataclass -class Measurement: - measurement_dir: Path - output_dir: Path = None - elements: Dict[str, Segment] = None - segments: Dict[str, Segment] = None - model_dir: Path = None - accel: str = None - year: str = None - ring: int = None - - def __str__(self): - return str(self.measurement_dir) - +from omc3_gui.segment_by_segment.measurement_model import OpticsMeasurement @dataclass @@ -32,8 +18,8 @@ class Settings: pass - class ItemDictModel: + """ Mixin-Class for a class that has a dictionary of items. """ def __init__(self): self.items = {} @@ -42,8 +28,9 @@ def try_emit(self, emit: bool = True): if not emit: return - if hasattr(self, "layoutChanged"): - self.layoutChanged.emit() + if hasattr(self, "dataChanged"): + # TODO: return which data has actually changed? + self.dataChanged.emit(self.index(0), self.index(len(self.items)-1), [Qt.EditRole]) def update_item(self, item): self.items[str(item)] = item @@ -84,24 +71,49 @@ def get_item_at(self, index: int) -> Any: return list(self.items.values())[index] - class MeasurementListModel(QtCore.QAbstractListModel, ItemDictModel): - items: Dict[str, Measurement] # for the IDE + items: Dict[str, OpticsMeasurement] # for the IDE + + class ColorRoles(enum.IntEnum): + NONE = 0 + BEAM1 = enum.auto() + BEAM2 = enum.auto() + RING1 = enum.auto() + RING2 = enum.auto() + RING3 = enum.auto() + RING4 = enum.auto() + + @classmethod + def get_color(cls, meas: OpticsMeasurement) -> int: + if meas.accel == "lhc": + return getattr(cls, f"BEAM{meas.beam}") + + if meas.accel == "psb": + return getattr(cls, f"RING{meas.ring}") + + return cls.NONE def __init__(self, *args, **kwargs): super(QtCore.QAbstractListModel, self).__init__(*args, **kwargs) super(ItemDictModel, self).__init__() def data(self, index: QtCore.QModelIndex, role: int = Qt.DisplayRole): - meas: Measurement = self.get_item_at(index.row()) + + meas: OpticsMeasurement = self.get_item_at(index.row()) if role == Qt.DisplayRole: # https://doc.qt.io/qt-5/qt.html#ItemDataRole-enum - return str(meas) - + return meas.display() + + if role == Qt.ToolTipRole: + return meas.tooltip() + + if role == Qt.TextColorRole: + return self.ColorRoles.get_color(meas) + if role == Qt.EditRole: return meas - def rowCount(self, index): + def rowCount(self, index: QtCore.QModelIndex = None): return len(self.items) @@ -128,7 +140,7 @@ def rowCount(self, parent=QtCore.QModelIndex()): def columnCount(self, parent=QtCore.QModelIndex()): return len(self._COLUMNS) - def data(self, index, role=QtCore.Qt.DisplayRole): + def data(self, index: QtCore.QModelIndex, role=QtCore.Qt.DisplayRole): i = index.row() j = index.column() segment: Segment = self.get_item_at(i) diff --git a/omc3_gui/segment_by_segment/view.py b/omc3_gui/segment_by_segment/view.py index e1004d6..3814646 100644 --- a/omc3_gui/segment_by_segment/view.py +++ b/omc3_gui/segment_by_segment/view.py @@ -7,11 +7,12 @@ import pyqtgraph as pg from accwidgets.graph import StaticPlotWidget from accwidgets.graph.widgets.plotitem import ExViewBox -from qtpy import QtWidgets, uic +from qtpy import QtWidgets, uic, QtGui from qtpy.QtCore import Qt, Signal, Slot, QModelIndex from omc3_gui.plotting.classes import DualPlot -from omc3_gui.segment_by_segment.model import Measurement, MeasurementListModel, SegmentTableModel +from omc3_gui.segment_by_segment.model import MeasurementListModel, SegmentTableModel +from omc3_gui.segment_by_segment.measurement_model import OpticsMeasurement from omc3_gui.utils.base_classes import View from omc3_gui.utils.counter import HorizontalGridLayoutFiller from omc3_gui.utils.widgets import EditButton, OpenButton, RemoveButton, RunButton @@ -25,37 +26,41 @@ class Tab: class SbSWindow(View): WINDOW_TITLE = "OMC Segment-by-Segment" + # QtSignals need to be defined as class-attributes sig_load_button_clicked = Signal() sig_remove_button_clicked = Signal() sig_matcher_button_clicked = Signal() - sig_edit_measurement_button_clicked = Signal(Measurement) - sig_list_optics_double_clicked = Signal(Measurement) - - _tabs: Dict[str, DualPlot] - _tabs_widget: QtWidgets.QTabWidget - _list_view_measurements: QtWidgets.QListView - _table_segments: QtWidgets.QTableView - - # Buttons --- - button_load: QtWidgets.QPushButton - button_remove: QtWidgets.QPushButton - button_edit: QtWidgets.QPushButton - button_matcher: QtWidgets.QPushButton - - button_run_segment: QtWidgets.QPushButton - button_remove_segment: QtWidgets.QPushButton - button_copy_segment: QtWidgets.QPushButton - button_new_segment: QtWidgets.QPushButton - + sig_edit_measurement_button_clicked = Signal(OpticsMeasurement) + sig_list_optics_double_clicked = Signal(OpticsMeasurement) + def __init__(self, parent=None): super().__init__(parent) + + # List of UI elements accessible as instance-attributes: + # Widgets --- + self._tabs_widget: QtWidgets.QTabWidget = None + self._tabs: Dict[str, DualPlot] = None + self._list_view_measurements: QtWidgets.QListView = None + self._table_segments: QtWidgets.QTableView = None + + # Buttons --- + self._button_load: QtWidgets.QPushButton = None + self._button_remove: QtWidgets.QPushButton = None + self._button_edit: QtWidgets.QPushButton = None + self._button_matcher: QtWidgets.QPushButton = None + + self._button_run_segment: QtWidgets.QPushButton = None + self._button_remove_segment: QtWidgets.QPushButton = None + self._button_copy_segment: QtWidgets.QPushButton = None + self._button_new_segment: QtWidgets.QPushButton = None + self._build_gui() - self.connect_signals() + self._connect_signals() self.plot() - def connect_signals(self): - self.button_load.clicked.connect(self._handle_load_files_button_clicked) + def _connect_signals(self): + self._button_load.clicked.connect(self._handle_load_files_button_clicked) self._list_view_measurements.doubleClicked.connect(self._handle_list_measurements_double_clicked) @Slot() @@ -82,9 +87,7 @@ def build_navigation_top(): nav_top.setLayout(layout) layout.addWidget(QtWidgets.QLabel("Loaded Optics:")) - self._list_view_measurements = QtWidgets.QListView() - self._list_view_measurements.setModel(MeasurementListModel()) - self._list_view_measurements.setSelectionMode(QtWidgets.QAbstractItemView.ExtendedSelection) + self._list_view_measurements = MeasurementListView() layout.addWidget(self._list_view_measurements) def build_measurement_buttons(): @@ -93,19 +96,19 @@ def build_measurement_buttons(): load = OpenButton("Load") grid_buttons_filler.add(load) - self.button_load = load + self._button_load = load remove = RemoveButton() grid_buttons_filler.add(remove) - self.button_remove = remove + self._button_remove = remove edit = EditButton() grid_buttons_filler.add(edit, col_span=2) - self.button_edit = edit + self._button_edit = edit matcher = RunButton("Run Matcher") grid_buttons_filler.add(matcher, col_span=2) - self.button_matcher = matcher + self._button_matcher = matcher return grid_buttons @@ -138,19 +141,19 @@ def build_segment_buttons(): run = RunButton("Run Segment(s)") grid_buttons_filler.add(run, col_span=3) - self.button_run_segment = run + self._button_run_segment = run new = OpenButton("New") grid_buttons_filler.add(new) - self.button_new_segment = new + self._button_new_segment = new copy = EditButton("Copy") grid_buttons_filler.add(copy) - self.button_copy_segment = copy + self._button_copy_segment = copy remove = RemoveButton("Remove") grid_buttons_filler.add(remove) - self.button_remove_segment = remove + self._button_remove_segment = remove return grid_buttons layout.addLayout(build_segment_buttons()) @@ -207,3 +210,39 @@ def clicked(self, item, points, ev): print('clicked') +class MeasurementListView(QtWidgets.QListView): + + def __init__(self): + super().__init__() + self.setModel(MeasurementListModel()) + self.setItemDelegate(ColoredItemDelegate()) + self.setSelectionMode(QtWidgets.QAbstractItemView.ExtendedSelection) + tooltip_style = """ + QToolTip { + background-color: #F0F0F0; /* Light gray background */ + color: #333333; /* Dark gray text */ + border: 1px solid #808080; /* Gray border */ + font-family: "Courier New", monospace; /* Monospaced font */ + } + """ + self.setStyleSheet(tooltip_style) + + +class ColoredItemDelegate(QtWidgets.QStyledItemDelegate): + + COLOR_MAP = { + MeasurementListModel.ColorRoles.NONE: "#000000", + MeasurementListModel.ColorRoles.BEAM1: "#0000ff", + MeasurementListModel.ColorRoles.BEAM2: "#ff0000", + # todo: what are the PSB ring colors? + MeasurementListModel.ColorRoles.RING1: "#4CAF50", + MeasurementListModel.ColorRoles.RING2: "#FF9800", + MeasurementListModel.ColorRoles.RING3: "#673AB7", + MeasurementListModel.ColorRoles.RING4: "#E91E63", + } + def paint(self, painter, option, index): + # Customize the text color + color = self.COLOR_MAP[index.data(Qt.TextColorRole)] + option.palette.setColor(QtGui.QPalette.Text, QtGui.QColor(color)) + + super().paint(painter, option, index) \ No newline at end of file diff --git a/omc3_gui/utils/dialogs.py b/omc3_gui/utils/file_dialogs.py similarity index 82% rename from omc3_gui/utils/dialogs.py rename to omc3_gui/utils/file_dialogs.py index 5c72f0d..bdfb454 100644 --- a/omc3_gui/utils/dialogs.py +++ b/omc3_gui/utils/file_dialogs.py @@ -8,7 +8,9 @@ LOGGER = logging.getLogger(__name__) +# Open Dialog Windows ---------------------------------------------------------- class OpenAnyFileDialog(QFileDialog): + def __init__(self, **kwargs) -> None: """ Quick dialog to open any kind of file. Modifies QFileDialog, and allows only kwargs to be passed. @@ -51,3 +53,10 @@ def __init__(self, caption = "Select Folder", **kwargs) -> None: def run_selection_dialog(self) -> Path: return super().run_selection_dialog()[0] + +class OpenFilesDialog(OpenAnyFileDialog): + def __init__(self, caption = "Select Files", **kwargs) -> None: + super().__init__(caption=caption, **kwargs) # parent, directory, filter, options + self.setFileMode(QFileDialog.ExistingFiles) + icon = QApplication.style().standardIcon(QStyle.SP_FileIcon) + self.setWindowIcon(icon)