feat: add selector regions for data plots

q0_gui.py

@@ -1,14 +1,16 @@
 from functools import partial
 from typing import Dict, Optional
 
+import numpy as np
 from PyQt5.QtCore import pyqtSignal, pyqtSlot
-from PyQt5.QtWidgets import QHBoxLayout, QVBoxLayout
+from PyQt5.QtWidgets import QHBoxLayout, QVBoxLayout, QPushButton
 from lcls_tools.common.frontend.display.util import showDisplay
 from lcls_tools.superconducting.sc_linac_utils import ALL_CRYOMODULES
 from pydm import Display
-from pyqtgraph import PlotWidget, plot
+from pyqtgraph import LinearRegionItem, PlotWidget, plot
 
 import q0_gui_utils
+import q0_utils
 from q0_gui_utils import CalibrationWorker
 from q0_linac import Q0Cryomodule, Q0_CRYOMODULES
 from q0_utils import ValveParams
@@ -24,6 +26,8 @@ def ui_filename(self):
     def __init__(self, parent=None, args=None):
         super().__init__(parent=parent, args=args)
 
+        self.region_button = None
+        self.region_button_text = ""
         self.selectedCM: Optional[Q0Cryomodule] = None
         self.ui.cm_combobox.addItems([""] + ALL_CRYOMODULES)
         self.ui.cm_combobox.currentTextChanged.connect(self.update_cm)
@@ -123,19 +127,55 @@ def update_cm(self, current_text):
             for cavity in self.selectedCM.cavities.values():
                 self.cav_amp_controls[cavity.number].connect(cavity)
 
+    def show_hide_regions(self, runs, measurement_plot, plot_items, region_event=None):
+        if any(isinstance(region_selector, LinearRegionItem) for region_selector in plot_items):
+            self.region_button_text = 'Show Measurement Regions'
+            self.region_button.setText(self.region_button_text)
+
+            for run in runs:
+                region_item = run.region
+                measurement_plot.removeItem(region_item)
+                plot_items.remove(region_item)
+        else:
+            self.region_button_text = 'Hide Measurement Regions'
+            self.region_button.setText(self.region_button_text)
+            for heater_run in runs:
+                linear_item = heater_run.region
+                if region_event:
+                    linear_item.sigRegionChangeFinished.connect(region_event)
+                measurement_plot.addItem(linear_item)
+                plot_items.append(linear_item)
+
+    def q0_update(self, region_item):
+        if region_item == self.selectedCM.q0_measurement.rf_run.region:
+            print("Updated Q0 -------------")
+            _ = self.selectedCM.q0_measurement.q0
+
     @pyqtSlot()
     def show_q0_data(self):
         if not self.q0_window:
             self.q0_window = Display()
             self.q0_window.setWindowTitle("Q0 Plots")
-            layout: QHBoxLayout = QHBoxLayout()
+            complete_layout: QVBoxLayout = QVBoxLayout()
+            plot_layout: QHBoxLayout = QHBoxLayout()
             self.q0_data_plot: PlotWidget = plot()
             self.q0_data_plot.setTitle("Q0 Data")
             self.q0_fit_plot: PlotWidget = plot()
+            self.region_button: QPushButton = QPushButton()
             self.q0_fit_plot.setTitle("Heat On Calibration Curve (with adjustments)")
-            layout.addWidget(self.q0_data_plot)
-            layout.addWidget(self.q0_fit_plot)
-            self.q0_window.setLayout(layout)
+            self.region_button.setMaximumWidth(250)
+            plot_layout.addWidget(self.q0_data_plot)
+            plot_layout.addWidget(self.q0_fit_plot)
+            complete_layout.addWidget(self.region_button)
+            complete_layout.addLayout(plot_layout)
+            self.q0_window.setLayout(complete_layout)
+            self.region_button.clicked.connect(
+                lambda: self.show_hide_regions([measurement.heater_run, measurement.rf_run],
+                                               self.q0_data_plot,
+                                               self.q0_data_plot_items, region_event=self.q0_update))
+
+        self.region_button_text = 'Show Measurement Regions'
+        self.region_button.setText(self.region_button_text)
 
         while self.q0_data_plot_items:
             self.q0_data_plot.removeItem(self.q0_data_plot_items.pop())
@@ -144,16 +184,17 @@ def show_q0_data(self):
             self.q0_fit_plot.removeItem(self.q0_fit_plot_items.pop())
 
         measurement = self.selectedCM.q0_measurement
+
         self.q0_data_plot_items.append(
             self.q0_data_plot.plot(
-                list(measurement.rf_run.ll_data.keys()),
-                list(measurement.rf_run.ll_data.values()),
+                list(measurement.rf_run.complete_ll_data().keys()),
+                list(measurement.rf_run.complete_ll_data().values()),
             )
         )
         self.q0_data_plot_items.append(
             self.q0_data_plot.plot(
-                list(measurement.heater_run.ll_data.keys()),
-                list(measurement.heater_run.ll_data.values()),
+                list(measurement.heater_run.complete_ll_data().keys()),
+                list(measurement.heater_run.complete_ll_data().values()),
             )
         )
 
@@ -176,19 +217,35 @@ def show_q0_data(self):
 
         showDisplay(self.q0_window)
 
+    def calibration_update(self):
+        print(f"Updated dLL/dt/dheat: {self.selectedCM.calibration.dLLdt_dheat}")
+
     @pyqtSlot()
     def show_calibration_data(self):
         if not self.calibration_window:
             self.calibration_window = Display()
             self.calibration_window.setWindowTitle("Calibration Plots")
-            layout: QHBoxLayout = QHBoxLayout()
+            complete_layout: QVBoxLayout = QVBoxLayout()
+            plot_layout: QHBoxLayout = QHBoxLayout()
             self.calibration_data_plot: PlotWidget = plot()
             self.calibration_data_plot.setTitle("Calibration Data")
             self.calibration_fit_plot: PlotWidget = plot()
             self.calibration_fit_plot.setTitle("Heat vs dll/dt")
-            layout.addWidget(self.calibration_data_plot)
-            layout.addWidget(self.calibration_fit_plot)
-            self.calibration_window.setLayout(layout)
+            self.region_button: QPushButton = QPushButton()
+            self.region_button.setMaximumWidth(250)
+            complete_layout.addWidget(self.region_button)
+            plot_layout.addWidget(self.calibration_data_plot)
+            plot_layout.addWidget(self.calibration_fit_plot)
+            complete_layout.addLayout(plot_layout)
+            self.calibration_window.setLayout(complete_layout)
+            self.region_button.clicked.connect(lambda: self.show_hide_regions(self.selectedCM.calibration_buffer_plots[
+                                                                                  self.selectedCM.calibration.time_stamp
+                                                                              ].heater_runs,
+                                                                              self.calibration_data_plot,
+                                                                              self.calibration_data_plot_items,
+                                                                              region_event=self.calibration_update))
+        self.region_button_text = 'Show Measurement Regions'
+        self.region_button.setText(self.region_button_text)
 
         while self.calibration_data_plot_items:
             self.calibration_data_plot.removeItem(
@@ -199,14 +256,18 @@ def show_calibration_data(self):
             self.calibration_fit_plot.removeItem(self.calibration_fit_plot_items.pop())
 
         dll_dts = []
+        if self.selectedCM.calibration.time_stamp not in self.selectedCM.calibration_buffer_plots:
+            self.selectedCM.calibration_buffer_plots[self.selectedCM.calibration.time_stamp] = self.selectedCM.calibration
 
-        for heater_run in self.selectedCM.calibration.heater_runs:
+        for heater_run in self.selectedCM.calibration_buffer_plots[self.selectedCM.calibration.time_stamp].heater_runs:
             self.calibration_data_plot_items.append(
                 self.calibration_data_plot.plot(
-                    list(heater_run.ll_data.keys()), list(heater_run.ll_data.values())
+                    list(heater_run.complete_ll_data().keys()), list(heater_run.complete_ll_data().values())
                 )
             )
+
             dll_dts.append(heater_run.dll_dt)
+
             self.calibration_fit_plot_items.append(
                 self.calibration_fit_plot.plot(
                     [heater_run.average_heat], [heater_run.dll_dt], pen=None, symbol="o"

q0_linac.py

@@ -19,6 +19,7 @@
     StepperTuner,
 )
 from numpy import floor, linspace, sign
+from pyqtgraph import LinearRegionItem
 from scipy.signal import medfilt
 from scipy.stats import linregress
 
@@ -42,7 +43,6 @@ def load_data(self):
         with open(self.cryomodule.calib_data_file, "r+") as f:
             all_data: Dict = json.load(f)
             data: Dict = all_data[self.time_stamp]
-
             for heater_run_data in data.values():
                 run = q0_utils.HeaterRun(heater_run_data["Desired Heat Load"])
                 run._start_time = datetime.strptime(
@@ -59,6 +59,11 @@ def load_data(self):
 
                 run.ll_data = ll_data
                 run.average_heat = heater_run_data[q0_utils.JSON_HEATER_READBACK_KEY]
+
+                times = np.array(list(run.ll_data.keys()))
+                run.region = LinearRegionItem(values=[np.min(times), np.max(times)],
+                                              bounds=[np.min(times), np.max(times)],
+                                              swapMode='block')
 
                 self.heater_runs.append(run)
 
@@ -107,23 +112,25 @@ def save_results(self):
 
     @property
     def dLLdt_dheat(self):
-        if not self._slope:
-            heat_loads = []
-            dll_dts = []
-            for run in self.heater_runs:
-                heat_loads.append(run.average_heat)
-                dll_dts.append(run.dll_dt)
-
-            slope, intercept, r_val, p_val, std_err = linregress(heat_loads, dll_dts)
+        heat_loads = []
+        dll_dts = []
+        for run in self.heater_runs:
+            heat_loads.append(run.average_heat)
+            dll_dts.append(run.dll_dt)
 
-            if np.isnan(slope):
-                self._slope = None
-            else:
-                self.adjustment = intercept
-                self._slope = slope
+        slope, intercept, r_val, p_val, std_err = linregress(heat_loads, dll_dts)
 
+        if np.isnan(slope):
+            self._slope = None
+        else:
+            self.adjustment = intercept
+            self._slope = slope
         return self._slope
 
+    @dLLdt_dheat.setter
+    def dLLdt_dheat(self, value: float):
+        self._slope = value
+
     def get_heat(self, dll_dt: float):
         return (dll_dt - self.adjustment) / self.dLLdt_dheat
 
@@ -235,6 +242,11 @@ def load_data(self, time_stamp: str):
 
             self.rf_run.avg_pressure = rf_run_data[q0_utils.JSON_AVG_PRESS_KEY]
 
+            for run in [self.heater_run, self.rf_run]:
+                times = np.array(list(run.ll_data.keys()))
+                run.region = LinearRegionItem(values=[np.min(times), np.max(times)],
+                                              bounds=[np.min(times), np.max(times)],
+                                              swapMode='block')
         self.save_data()
 
     def save_data(self):
@@ -281,8 +293,7 @@ def save_results(self):
 
     @property
     def raw_heat(self):
-        if not self._raw_heat:
-            self._raw_heat = self.cryomodule.calibration.get_heat(self.rf_run.dll_dt)
+        self._raw_heat = self.cryomodule.calibration.get_heat(self.rf_run.dll_dt)
         return self._raw_heat
 
     @property
@@ -296,28 +307,29 @@ def adjustment(self):
 
     @property
     def heat_load(self):
-        if not self._heat_load:
-            self._heat_load = self.raw_heat + self.adjustment
+        self._heat_load = self.raw_heat + self.adjustment
         return self._heat_load
 
     @property
     def q0(self):
-        if not self._q0:
-            sum_square_amp = 0
+        sum_square_amp = 0
 
-            for amp in self.rf_run.amplitudes.values():
-                sum_square_amp += amp**2
+        for amp in self.rf_run.amplitudes.values():
+            sum_square_amp += amp**2
 
-            effective_amplitude = np.sqrt(sum_square_amp)
+        effective_amplitude = np.sqrt(sum_square_amp)
 
-            self._q0 = q0_utils.calc_q0(
-                amplitude=effective_amplitude,
-                rf_heat_load=self.heat_load,
-                avg_pressure=self.rf_run.avg_pressure,
-                cav_length=self.cryomodule.cavities[1].length,
-            )
+        self._q0 = q0_utils.calc_q0(
+            amplitude=effective_amplitude,
+            rf_heat_load=self.heat_load,
+            avg_pressure=self.rf_run.avg_pressure,
+            cav_length=self.cryomodule.cavities[1].length,
+        )
         return self._q0
 
+    @q0.setter
+    def q0(self, value: float):
+        self._q0 = value
 
 class Q0Cavity(Cavity):
     def __init__(
@@ -396,6 +408,9 @@ def __init__(
 
         self._ds_level_pv_obj: Optional[PV] = None
 
+        self.calibration_buffer_plots: Dict[str, Calibration] = {}
+        self.q0_measurement_buffer_plots: Dict[str, Q0Measurement] = {}
+
     def __str__(self):
         return f"CM{self.name}"
 
@@ -469,8 +484,8 @@ def shut_off(self):
         caput(self.jtAutoSelectPV, 1, wait=True)
         print("Turning cavities and SSAs off")
         for cavity in self.cavities.values():
-            cavity.turnOff()
-            cavity.ssa.turnOff()
+            cavity.turn_off()
+            cavity.ssa.turn_off()
 
     @property
     def heater_power(self):
@@ -510,7 +525,7 @@ def fill(self, desired_level=q0_utils.MAX_DS_LL, turn_cavities_off: bool = True)
 
         if turn_cavities_off:
             for cavity in self.cavities.values():
-                cavity.turnOff()
+                cavity.turn_off()
 
         self.waitForLL(desired_level)
 
@@ -565,7 +580,7 @@ def getRefValveParams(self, start_time: datetime, end_time: datetime):
                 # We only want to use time periods in which there were no
                 # changes made to the heater settings
                 if len(des_val_set) == 1:
-                    des_pos = round(np.mean(data.values[self.jtValveReadbackPV]), 1)
+                    des_pos = round(np.mean(data.values[self.jt_valve_readback_pv]), 1)
                     heater_des = des_val_set.pop()
                     heater_act = np.mean(data.values[self.heater_readback_pv])
 
@@ -717,7 +732,7 @@ def takeNewQ0Measurement(
         duration = (end_time - start_time).total_seconds() / 3600
         print("Duration in hours: {DUR}".format(DUR=duration))
 
-        print("Caluclated Q0: ", self.q0_measurement.q0)
+        print("Calculated Q0: ", self.q0_measurement.q0)
         self.q0_measurement.save_results()
         self.restore_cryo()
 
@@ -737,14 +752,22 @@ def setup_for_q0(
         self.fill(desired_ll)
 
     def load_calibration(self, time_stamp: str):
-        self.calibration: Calibration = Calibration(
-            time_stamp=time_stamp, cryomodule=self
-        )
-        self.calibration.load_data()
+        if time_stamp not in self.calibration_buffer_plots:
+            self.calibration: Calibration = Calibration(
+                time_stamp=time_stamp, cryomodule=self
+            )
+            self.calibration.load_data()
+            self.calibration_buffer_plots[time_stamp] = self.calibration
+        else:
+            self.calibration = self.calibration_buffer_plots[time_stamp]
 
     def load_q0_measurement(self, time_stamp):
-        self.q0_measurement: Q0Measurement = Q0Measurement(self)
-        self.q0_measurement.load_data(time_stamp)
+        if time_stamp not in self.q0_measurement_buffer_plots:
+            self.q0_measurement: Q0Measurement = Q0Measurement(self)
+            self.q0_measurement.load_data(time_stamp)
+            self.q0_measurement_buffer_plots[time_stamp] = self.q0_measurement
+        else:
+            self.q0_measurement = self.q0_measurement_buffer_plots[time_stamp]
 
     def takeNewCalibration(
         self,