pretend to load measurements, parse info, color and tooltips

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))
 
 
-

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

omc3_gui/segment_by_segment/model.py

@@ -1,39 +1,25 @@
-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
 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)