element lines and help

Author: JoschD

omc3_gui/plotting/element_lines.py

@@ -0,0 +1,53 @@
+""" 
+Element Line Plotter
+--------------------
+
+This module contains functions to plot element lines with pyqtgraph.
+"""
+import numpy as np
+import pandas as pd
+import tfs
+import pyqtgraph as pg
+import logging
+from omc3.optics_measurements.constants import S
+from qtpy.QtCore import Qt
+
+LOGGER = logging.getLogger(__name__)
+LENGTH: str = "LENGTH"
+
+def plot_element_lines(plot: pg.PlotWidget, data_frame: tfs.TfsDataFrame, ranges: list[tuple[str, str]], start_zero: bool): 
+    """ 
+    Plot vertical lines on the plot for elements in the data_frame in the given ranges.
+
+    Args:
+        plot (pg.PlotWidget): The plot to plot the lines into.
+        data_frame (tfs.TfsDataFrame): The data_frame to plot the lines from.
+        ranges (list[tuple[str, str]]): A list of tuples of the form (start_element, end_element).
+        start_zero (bool): Whether to start the plot from zero or not.  
+    """
+    if start_zero and not all(ranges[0][0] == r[0] for r in ranges):
+        LOGGER.warning("Not all ranges start at the same element. Using only the first!")
+        ranges = [ranges[0]]
+
+    plotItem = plot.plotItem
+
+    # Find start and end elements
+    start = min(ranges, key=lambda r: data_frame.loc[r[0], S])[0]
+    end = max(ranges, key=lambda r: data_frame.loc[r[1], S])[1]
+
+    # Select element range and do some wrapping gymnastics if needed
+    if data_frame.loc[start, S] <= data_frame.loc[end, S]:
+        s_elements = data_frame.loc[start:end, S]
+    else:
+        s_elements = data_frame.loc[start:, S] + data_frame.loc[:end, S]
+    
+    if start_zero:
+        s_elements = s_elements - s_elements.loc[start]
+        s_elements = s_elements + np.where(s_elements < 0, data_frame.headers[LENGTH], 0)
+
+    # Plot the lines - this takes a while, not sure how to improve
+    plotItem.disableAutoRange()  # speeds it up a bit
+    pen = pg.mkPen(color="grey", width=1, style=Qt.PenStyle.DotLine)
+    for element, x in s_elements.items():
+        plotItem.addLine(x=x, z=-10, pen=pen, label=element)
+

omc3_gui/segment_by_segment/help_view.py

@@ -0,0 +1,87 @@
+""" 
+Help Dialogs
+------------
+"""
+
+from qtpy.QtCore import Qt
+from qtpy.QtWidgets import QMessageBox
+
+def show_help_dialog():
+    help_text = """
+    <h3> Frequently Asked Questions </h3><br>
+    
+    <br>
+    <b> How do I open a Measurement ?</b><br>
+    
+    One way to open measurements automatically, is to give them as command line
+    arguments when starting the sbs_gui, either <i> -m </i> or <i> --measurements </i>.<br>
+    If you want to load them manually, you click the <i>Load</i> button.<br><br>
+
+    The given/selected measurements can be either <i>omc3 optics folders, folders containing sbs-json files
+    or sbs-json files</i> directly.
+    The latter are created automatically in the measurement-output folder when editing a loaded measurement.<br>
+
+    <br>
+    
+    <b> Do I have to invert my corrections when using them in the machine? </b><br>
+
+    YES! <i>(but it depends)</i><br>
+    The "corrections" here are actually used to match the model opttics to the 
+    mesured optics (see the info about the dashed "corr" line below). 
+    Therefore you have to invert them, to actually use them as corrections in the machine. <br>
+    <b> NOTE </b> that these are the MAD-X values. Make how the signs are actually
+    mapped in the machine! <br>
+
+    <br>
+
+    <b> What is the solid line ?</b><br>
+
+    The solid line is the difference between the Measurement and 
+    the propagated model, i.e. the Measurement at the start (or end) of the segment 
+    propagated through the nominal model via MAD-X.<br>
+    This line therefore shows you how much the optics deviate through the segment 
+    from the nominal model.<br> 
+
+    <br>   
+    
+    <b> What is the dashed line that says "corr" ?</b><br>
+
+    This is the difference between the <i>"corrected"</i> propagated model and the 
+    nominal propagated model.<br>
+    This means in both cases the measured values are used as initial conditions.
+    What you are trying to achieve is a match between the dashed and the solid line,
+    because that means that now your model matches the optics in the measured data.<br>
+
+    <br>
+    
+    <b> What is the dashed line that says "expct" ?</b><br>
+
+    This is the difference between the Measurement and the "corrected" propagated model
+    and is therefore the <i>expected</i> measured difference to the nominal model after 
+    applying the correction in the machine (same as in global correction).<br>
+    You can activate this view via the plot-settings <i>"Expectation"</i>.<br>
+
+    <br>
+
+    <b> Shortcuts </b><br>
+
+    In Graph:<br>
+    <i>Double-Click</i> : Zoom history back one step. <br>
+    <i>Right-Click</i> : Zoom history back one step. <br>
+    <i>Shift + Right-Click</i> : Zoom history back all steps. <br>
+    <i>Alt + Right-Click</i> : pyqtgraph context menu. <br>
+
+    <br>
+
+    In Measurements-List:<br>
+    <i>Double-Click</i> : Edit the Measurement.<br>
+
+    <br> 
+
+    """
+
+    msg_box = QMessageBox(icon=QMessageBox.Information)
+    msg_box.setWindowTitle("Help")
+    msg_box.setTextFormat(Qt.TextFormat.RichText)
+    msg_box.setText(help_text)
+    msg_box.exec_()

omc3_gui/segment_by_segment/main_controller.py

@@ -13,7 +13,6 @@
 from typing import TYPE_CHECKING
 
 from omc3.sbs_propagation import segment_by_segment
-from omc3.segment_by_segment.constants import corrections_madx
 from qtpy import QtWidgets
 from qtpy.QtCore import Slot
 
@@ -723,12 +722,6 @@ def load_segments_for_measurement(self, measurement: OpticsMeasurement):
             segment = SegmentDataModel(measurement, *segment_tuple)
             measurement.try_add_segment(segment)
 
-        for segment in measurement.segments:
-            corrections = measurement.output_dir / corrections_madx.format(segment.name)
-            if corrections.exists():
-                measurement.corrections = corrections
-                break  # for now they should all be the same corrections
-        
     @Slot()
     def save_segments(self):
         LOGGER.debug("Saving segments to a file.")
@@ -844,7 +837,7 @@ def plot(self, fail_ok: bool = True):
             log_function("Not plotting, no segments have been run.")
             return
 
-        if settings.same_start:
+        if settings.same_start and not settings.model_s:  # only an issue if they all start at 0
             starts = {re.sub(r"\.B[12]$", "", s.start, flags=re.IGNORECASE) for s in segments_data}
             if len(starts) > 1:
                 log_function("Not plotting, segments have different start BPMs (see 'Same Start' in settings).")

omc3_gui/segment_by_segment/main_view.py

@@ -22,6 +22,7 @@
 from qtpy.QtCore import QItemSelectionModel, QModelIndex, Qt, Signal, Slot
 
 from omc3_gui.plotting.classes import DualPlotWidget
+from omc3_gui.segment_by_segment.help_view import show_help_dialog
 from omc3_gui.segment_by_segment.main_model import (
     MeasurementListModel,
     SegmentTableModel,
@@ -153,16 +154,25 @@ def _add_menus(self):
 
         # insert before the last entry (which is "Exit")
         file_menu.insertAction(file_menu.actions()[-1], menu_settings)
-        
-        # Clear All ---
+
+        # Help --- 
         help_menu: QtWidgets.QMenu = self.get_action_by_title("Help")  # defined in View-class
+
+        # Clear All -
         menu_clear_all = QtWidgets.QAction("Reload Data", self)
         menu_clear_all.setIcon(
             QtWidgets.QApplication.style().standardIcon(QtWidgets.QStyle.StandardPixmap.SP_DialogResetButton)
         )
         menu_clear_all.triggered.connect(self.sig_menu_clear_all.emit)
         help_menu.insertAction(help_menu.actions()[-1], menu_clear_all)
 
+        menu_show_help = QtWidgets.QAction("Show Help", self)
+        menu_show_help.setIcon(
+            QtWidgets.QApplication.style().standardIcon(QtWidgets.QStyle.StandardPixmap.SP_MessageBoxQuestion)
+        )
+        menu_show_help.triggered.connect(show_help_dialog)  # more of a controller thing, but OK for this one
+        help_menu.insertAction(help_menu.actions()[-1], menu_show_help)
+
     def add_settings_to_menu(self, menu: str, settings: object, names: Sequence[str] | None = None, hook: callable = None):
         """ Add quick-access checkboxes to the menu which are connected to the respective attributes in settings.
         
@@ -433,13 +443,6 @@ def __init__(self):
         self.setSelectionMode(QtWidgets.QAbstractItemView.ExtendedSelection)
         self.setShowGrid(True)
         self.setStyleSheet(MONOSPACED_TOOLTIP)
-    
-    # def mousePressEvent(self, e: QtGui.QMouseEvent) -> None:
-    #     idx = self.indexAt(e.pos())
-    #     if e.button() == Qt.RightButton:
-    #         self.model().toggle_row(idx)  # rather a controller thing?
-    #         return 
-    #     super().mousePressEvent(e)
 
 
 class ColoredItemDelegate(QtWidgets.QStyledItemDelegate):

omc3_gui/segment_by_segment/measurement_model.py

@@ -21,6 +21,7 @@
     MODEL_DIRECTORY,
     PHASE_NAME,
 )
+from omc3.segment_by_segment.constants import corrections_madx
 from tfs.reader import read_headers
 
 from omc3_gui.ui_components.dataclass_ui import DirectoryPath, FilePath, metafield
@@ -228,16 +229,19 @@ def from_path(cls, path: Path) -> OpticsMeasurement:
                 LOGGER.error(f"JSON errror: {e!s}\nTrying to load as optics-measurement folder.")
 
         # Try to load from optics-measurement folder ---
-        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)
+            info["model_dir"] = model_dir
 
-        return cls(measurement_dir=path, model_dir=model_dir, **info)
+        if (path / corrections_madx).is_file():
+            info["corrections"] = path / corrections_madx
+        
+        return cls(measurement_dir=path, **info)
 
     @classmethod
     def from_json(cls, path: Path, measurement_dir: Path | None = None) -> OpticsMeasurement:

omc3_gui/segment_by_segment/plotting.py

@@ -7,7 +7,9 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
+from functools import cache
 import logging
+from pathlib import Path
 
 from omc3.definitions.optics import (
     S_COLUMN,
@@ -19,9 +21,13 @@
     ColumnsAndLabels,
 )
 from omc3.segment_by_segment.propagables import PropagableColumns
+from omc3.model.constants import TWISS_ELEMENTS_DAT
+from omc3.optics_measurements.constants import NAME
 from qtpy.QtCore import Qt
+import tfs
 
 from omc3_gui.plotting.classes import DualPlotWidget
+from omc3_gui.plotting.element_lines import plot_element_lines
 from omc3_gui.plotting.latex_to_html import latex_to_html_converter
 from omc3_gui.plotting.tfs_plotter import plot_dataframes
 from omc3_gui.segment_by_segment.segment_model import SegmentDataModel
@@ -173,8 +179,7 @@ def get_data(segment: SegmentDataModel, file_name: str):
     x_column = S_COLUMN
     if settings.model_s:
         x_column = S_MODEL_COLUMN
-    
-    
+
     # Loop over top/bottom plots ---
     for definition, plot in zip(definitions.plots, widget.plots):
         definition: PlotDefinition
@@ -189,6 +194,17 @@ def get_data(segment: SegmentDataModel, file_name: str):
             # continue anyway
             dataframes = {label: df for label, df in dataframes.items() if df is not None}
 
+        # Plot Model Elements ---
+        if settings.show_model:
+            model_dir = segments[0].measurement.model_dir
+            bpm_ranges = [(s.start, s.end) for s in segments]
+            plot_element_lines(
+                plot=plot,
+                data_frame=load_twiss_elements(model_dir),
+                ranges=bpm_ranges,
+                start_zero=not settings.model_s,
+            )
+
         # Loop over forward/backward plots ---
         for direction in ("forward", "backward"):
             if not getattr(settings, direction):  # user activated
@@ -218,3 +234,12 @@ def get_data(segment: SegmentDataModel, file_name: str):
         if settings.reset_zoom:
             plot.enableAutoRange()
 
+
+@cache
+def load_twiss_elements(model_dir: Path) -> tfs.TfsDataFrame:
+    """ Load the twiss elements from the model directory. 
+    Cache here, because that might take a moment, so better to keep the DataFrame in memory.
+    """
+    df = tfs.read_tfs(model_dir / TWISS_ELEMENTS_DAT, index=NAME)
+    df = df.loc[df.index.str.match(r"^(?!DRIFT).*"), [S_COLUMN.column]]  # we actually only need the s column and headers
+    return df

omc3_gui/segment_by_segment/settings.py

@@ -20,6 +20,7 @@ class MainSettings:
 
 @dataclass(slots=True)
 class PlotSettings:
+    show_model: bool = metafield("Show Model", "Show markers for the elements of the Model.", default=False)
     show_legend: bool = metafield("Show Legend", "Show legend.", default=True)
     marker_size: float = metafield("Marker Size", "Size of the markers.", default=8.5)
     expected: bool = metafield("Expectation", "Show expected value after correction instead of correction itself.", default=False)