better plotting and others

Author: JoschD

omc3_gui/plotting/classes.py

@@ -7,6 +7,7 @@
 import pyqtgraph as pg
 from accwidgets.graph import StaticPlotWidget
 from accwidgets.graph.widgets.plotitem import ExViewBox
+from accwidgets.graph.widgets.plotwidget import GridOrientationOptions
 from qtpy.QtCore import Signal
 
 
@@ -39,25 +40,31 @@ def __init__(self, *args, **kwargs) -> None:
     def plots(self) -> tuple[pg.PlotWidget, pg.PlotWidget]:
         return (self.top, self.bottom)
 
-    
     def clear(self) -> None:    
         for plot in self.plots:
             plot.clear()
 
+    def set_connect_x(self, connect: bool) -> None:
+        if connect:
+            self.top.setXLink(self.bottom)
+        else:
+            self.top.setXLink(None)
+            self.bottom.setXLink(None)
 
-    def connect_x(self) -> None:
-        pass
-
-    def connect_y(self) -> None:
-        pass
+    def set_connect_y(self, connect: bool) -> None:
+        if connect:
+            self.top.setYLink(self.bottom)
+        else:
+            self.top.setYLink(None)
+            self.bottom.setYLink(None)
 
 
 class PlotWidget(StaticPlotWidget):
 
     def __init__(self, *args, **kwargs) -> None:
         super().__init__(*args, **kwargs, viewBox=ZoomingViewBox())  # requires accwidgets >= 3.0.11
         self.setBackground("w")
-
+        self._set_show_grid(GridOrientationOptions.Both)
 
 class ZoomingViewBox(ExViewBox):
     """ ViewBox that imitates the bahavior of the Java-GUI a bit closer. 
@@ -68,6 +75,11 @@ def __init__(self, *args, **kwargs) -> None:
         super().__init__(*args, **kwargs)
         self.setMouseMode(ZoomingViewBox.RectMode)  # mode that makes zooming rectangles
 
+    def suggestPadding(self, axis):
+        if axis == 0:
+            return 0.0  # disable padding for x axis
+        return super().suggestPadding(axis)
+
     # def mouseDragEvent(self, ev):
 
     #     if ev.button() == QtCore.Qt.RightButton:

omc3_gui/segment_by_segment.py

@@ -10,7 +10,6 @@
  GUI:
  - Load segments from file or folder (check sbs/sbs_ files)
  - Save segments to file
- - Autoload segments when opening measurement folder
  - Pass measurement folders via cli args.
 
 Settings:
@@ -22,7 +21,6 @@
  - Going back through plot history on double-click
 
 """
-
 import sys
 from omc3_gui.segment_by_segment.main_controller import SbSController
 from omc3_gui.utils.log_handler import init_logging
@@ -34,4 +32,5 @@
 
 if __name__ == "__main__":
     init_logging()
-    sys.exit(SbSController.run_application(measurements=["/mnt/volume/jdilly/projects/omc3_gui/tst_SBStest_wACD/measured_optics"]))
+    # sys.exit(SbSController.run_application(measurements=["/mnt/volume/jdilly/projects/omc3_gui/tst_SBStest_wACD/measured_optics", "/afs/cern.ch/work/j/josch/temp.bbgui_output/2025-02-19/LHCB1/Results/omc3_B1_30cm_with_old_local_corrections/"]))
+    sys.exit(SbSController.run_application())

omc3_gui/segment_by_segment/defaults.py

@@ -4,11 +4,56 @@
 
 Defaults for segment by segment.
 """
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
 from omc3_gui.segment_by_segment.segment_model import SegmentTuple
 
+if TYPE_CHECKING: 
+    from omc3_gui.segment_by_segment.measurement_model import OpticsMeasurement
+
 DEFAULT_SEGMENTS =(
     SegmentTuple("IP1", "BPM.12L1", "BPM.12R1"),
     SegmentTuple("IP2", "BPM.12L2", "BPM.12R2"),
     SegmentTuple("IP5", "BPM.12L5", "BPM.12R5"),
     SegmentTuple("IP8", "BPM.12L8", "BPM.12R8"),
-)
+)
+
+
+LHC_ARCS = ('78', '81', '12', '23', '34', '45', '56', '67')  # 78 needs to be before 81 to make lr fit
+LHC_CORRECTORS = (
+    "kqf.a{}", 
+    "kqd.a{}", 
+    "kqt12.{}{}b{}", 
+    "kqtl11.{}{}b{}", 
+    "kq10.{}{}b{}",
+    "kq9.{}{}b{}",
+    "kq8.{}{}b{}",
+    "kq7.{}{}b{}",
+    "kq6.{}{}b{}",
+    "kq5.{}{}b{}",
+    "kq4.{}{}b{}",
+    "kqx.{}{}b{}",
+    "kqsx3.{}{}b{}",
+    "ktqx2.{}{}b{}",
+    "ktqx1.{}{}b{}",
+    "kqx.{}{}b{}",
+)
+
+def get_default_correctors(measurement: OpticsMeasurement) -> str:
+    text = ""
+    if measurement.accel == "lhc":
+        text = ""
+        for ip in (1, 2, 5, 8):
+            text += f"! IP{ip} -----\n"
+            arcs = [arc for arc in LHC_ARCS if str(ip) in arc]
+            for arc, side in zip(arcs, "lr"):
+                lhc_correctors = LHC_CORRECTORS if side == "l" else LHC_CORRECTORS[::-1]
+                for corrector in lhc_correctors:
+                    if "a" in corrector:
+                        corrector = corrector.format(arc)
+                    else:
+                        corrector = corrector.format(side, ip, measurement.beam)
+                    text += f"! {corrector} = {corrector};\n"
+            text += "\n"
+    return text 

omc3_gui/segment_by_segment/main_controller.py

@@ -12,10 +12,11 @@
 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
 
-from omc3_gui.segment_by_segment.defaults import DEFAULT_SEGMENTS
+from omc3_gui.segment_by_segment.defaults import DEFAULT_SEGMENTS, get_default_correctors
 from omc3_gui.segment_by_segment.plotting import plot_segment_data
 from omc3_gui.segment_by_segment.settings import PlotSettings, Settings
 from omc3_gui.segment_by_segment.main_model import SegmentTableModel
@@ -70,7 +71,7 @@ def connect_signals(self):
         view.add_settings_to_menu(
             menu="View",
             settings=self.settings.plotting,
-            hook=partial(self.plot, weak=True),
+            hook=partial(self.plot, weak=True),  # update plots if possible
         )
 
         # Measurements -------------------------------------------------------------
@@ -421,6 +422,11 @@ def edit_corrections(self) -> None:
         # Open the TextEditor for the selected correction file ---
         LOGGER.debug(f"Opening TextEditor for correction file: {correction_file}.")
         edit_dialog = TextEditorDialog(correction_file)
+
+        if not correction_file.exists() and self.settings.main.suggest_correctors:
+            text = get_default_correctors(selected_measurements[0])  # bit hacky but ok for now?
+            edit_dialog.text_edit.setPlainText(text)
+
         edit_dialog.exec_()
 
 
@@ -659,6 +665,12 @@ 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.")
@@ -726,26 +738,30 @@ def clear_all():
             )
 
             # For Real Use: Run Task ---
-            # LOGGER.info(f"Starting {measurement_task.message}")
-            # self._add_running_task(task=measurement_task)
-            # measurement_task.start()
+            LOGGER.info(f"Starting {measurement_task.message}")
+            self._add_running_task(task=measurement_task)
+            measurement_task.start()
 
             # For Debugging: Start sbs directly ---
-            sbs_function()
-            clear_all()
-            LOGGER.info(f"Finished {measurement_task.message}")
+            # sbs_function()
+            # clear_all()
+            # LOGGER.info(f"Finished {measurement_task.message}")
             # -------------------------------------
 
 # Plotting ---------------------------------------------------------------------
     def plot(self, weak: bool = False):
         """ Trigger a plot update with the currently selected segments. """
         view: SbSWindow = self._view
         settings: PlotSettings = self.settings.plotting
+        definition, widget = view.get_current_tab()
 
         if not settings.forward and not settings.backward:
             LOGGER.error("Please enable at least one propagation method to show.")
             return
 
+        widget.set_connect_x(settings.connect_x)
+        widget.set_connect_y(settings.connect_y)
+
         segments = view.get_selected_segments()
         if len(segments) != 1:
             if not weak:
@@ -755,7 +771,6 @@ def plot(self, weak: bool = False):
         self.clear_plots()
 
         segments_data: list[SegmentDataModel] = segments[0].segments
-        definition, widget = view.get_current_tab()
         plot_segment_data(
             widget=widget, 
             definition=definition, 

omc3_gui/segment_by_segment/main_view.py

@@ -144,6 +144,7 @@ def update_settings(value: bool, name: str):
             if hook is not None:
                 hook()
 
+        qmenu.addSeparator()
         for field in fields(settings): 
             if names is not None and field.name not in names:
                 continue
@@ -154,9 +155,9 @@ def update_settings(value: bool, name: str):
             entry = QtWidgets.QAction(label, self)
             entry.setCheckable(True)
             entry.setChecked(getattr(settings, field.name))
-
             entry.toggled.connect(partial(update_settings, name=field.name))
             qmenu.addAction(entry)
+        qmenu.addSeparator()
 
     def update_menu_settings(self, menu: str, settings: object, names: Sequence[str] | None = None):
         """ Update the menu settings. 

omc3_gui/segment_by_segment/plotting.py

@@ -26,6 +26,7 @@ def plot_segment_data(widget: DualPlot, definition: ColumnsAndLabels, segments:
     Plot the given segments with the given definition. 
     """
     s_column = S_COLUMN
+    
     for plane, plot in zip("xy", [widget.top, widget.bottom]): 
         data_name = f"{definition.text_label}_{plane}"  # coincides with the name in TfsCollection
 

omc3_gui/segment_by_segment/settings.py

@@ -7,7 +7,7 @@
 from __future__ import annotations
 
 from pathlib import Path
-from dataclasses import dataclass, field, fields
+from dataclasses import dataclass, field
 
 from omc3_gui.ui_components.dataclass_ui import metafield
 
@@ -16,16 +16,17 @@ class MainSettings:
     cwd: Path = metafield("Working Directory", "Current working directory. Used for default path when opening file selection dialogs.", default=Path.cwd())
     autoload_segments: bool = metafield("Autoload Segments", "Automatically try to load existing segments when loading a measurement.", default=True) 
     autodefault_segments: bool = metafield("Auto-Add Default Segments", "Automatically add default segments when loading a measurement.", default=False) 
+    suggest_correctors: bool = metafield("Suggest Correctors", "Suggest correctors when editing a new correction file.", default=True)
 
 
 @dataclass(slots=True)
 class PlotSettings:
-    # connect_x: bool = metafield("Connect X", "Connect X axes.", default=False)  # TODO
-    # connect_y: bool = metafield("Connect Y", "Connect Y axes.", default=False)  # TODO
     show_legend: bool = metafield("Show Legend", "Show legend.", default=True)
-    forward: bool = metafield("Forward Propagation", "Show forward propagation.", default=True)
-    backward: bool = metafield("Backward Propagation", "Show backward propagation.", default=True)
     expected: bool = metafield("Expectation", "Show expected value after correction instead of correction itself.", default=False)
+    forward: bool = metafield("Forward Propagation", "Show forward propagation.", default=True)
+    backward: bool = metafield("Backward Propagation", "Show backward propagation.", default=False)
+    connect_x: bool = metafield("Connect X", "Connect X axes.", default=True)
+    connect_y: bool = metafield("Connect Y", "Connect Y axes.", default=False)
 
 
 @dataclass(slots=True)

omc3_gui/ui_components/dataclass_ui/controller.py

@@ -376,6 +376,8 @@ def set_value(value):
 
     else:   # Any kind of widget should be able to handle strings.      
         def get_value():
+            if not widget.text():
+                return None
             return field_type(widget.text())
 
         def set_value(value):