Plotting Windows (#416)

* New Classes for GUI-Windows to organize plots
* Plotting from plot_correction_test now in window
* Plotting from plot_tfs now in window
* bumpversion, changelog
* save-opt only if output-dir is set

Author: JoschD

CHANGELOG.md

@@ -1,6 +1,13 @@
 # OMC3 Changelog
 
+#### 2023-06-05 - v0.11.0 - _jdilly_
+
+- Added:
+ - `omc3.plotting.utils.windows`: Qt-based windows and widgets for matplotlib-figure organization.
+ - Using the new windows in `omc3.plotting.plot_checked_corrections` and `omc3.plotting.plot_tfs`
+
 #### 2023-05-15 - v0.10.0 - _jdilly_
+
 - Added:
   - `omc3.check_corrections`: A new feature to check the validity of corrections.
   - `omc3.plotting.plot_checked_corrections`: Function to plot the checked corrections.

doc/modules/plotting.rst

@@ -54,3 +54,6 @@ Common Plot Functionality
     :members:
     :noindex:
 
+.. automodule:: omc3.plotting.utils.windows
+    :members:
+    :noindex:

omc3/__init__.py

@@ -11,7 +11,7 @@
 __title__ = "omc3"
 __description__ = "An accelerator physics tools package for the OMC team at CERN."
 __url__ = "https://github.com/pylhc/omc3"
-__version__ = "0.10.0"
+__version__ = "0.11.0"
 __author__ = "pylhc"
 __author_email__ = "[email protected]"
 __license__ = "MIT"

omc3/plotting/optics_measurements/utils.py

@@ -10,6 +10,7 @@
 from typing import Iterable, Optional
 
 from matplotlib import pyplot as plt
+from matplotlib.figure import Figure
 from numpy.typing import ArrayLike
 
 
@@ -33,8 +34,9 @@ class DataSet:
 class FigureContainer:
     """Container for attaching additional information to one figure."""
     def __init__(self, id_: str, path: Path, axes_ids: Iterable[str]) -> None:
-        self.fig, axs = plt.subplots(nrows=len(axes_ids))
-        self.fig.canvas.manager.set_window_title(id_)
+        self.fig = Figure()
+        axs = self.fig.subplots(nrows=len(axes_ids))
+        self.id = id_
 
         if len(axes_ids) == 1:
             axs = [axs]
@@ -55,6 +57,9 @@ class FigureCollector:
     def __init__(self) -> None:
         self.fig_dict = OrderedDict()   # dictionary of matplotlib figures, for output
         self.figs = OrderedDict()       # dictionary of FigureContainers, used internally
+    
+    def __len__(self) -> int:
+        return len(self.figs)
 
     def add_data_for_id(self, figure_id: str, label: str, data: DataSet,
                         x_label: str, y_label: str,
@@ -93,4 +98,3 @@ def safe_format(label: str, insert: str) -> Optional[str]:
         return label
     except AttributeError:  # label is None
         return None
-

omc3/plotting/plot_checked_corrections.py

@@ -152,29 +152,43 @@
 
 
 """
+import re
 from pathlib import Path
-from typing import Dict, Iterable, Set
+from typing import Dict, Iterable, Set, Union
 
 from matplotlib import pyplot as plt
+from matplotlib.figure import Figure
 
 import tfs
 from generic_parser import DotDict, EntryPointParameters, entrypoint
-from omc3.correction.constants import (CORRECTED_LABEL, UNCORRECTED_LABEL, CORRECTION_LABEL, EXPECTED_LABEL,
-                                       COUPLING_NAME_TO_MODEL_COLUMN_SUFFIX)
-from omc3.definitions.optics import FILE_COLUMN_MAPPING, ColumnsAndLabels, RDT_COLUMN_MAPPING
+from omc3.correction.constants import (
+    CORRECTED_LABEL, UNCORRECTED_LABEL, CORRECTION_LABEL, EXPECTED_LABEL,
+    COUPLING_NAME_TO_MODEL_COLUMN_SUFFIX
+)
+from omc3.definitions.optics import (
+    FILE_COLUMN_MAPPING, ColumnsAndLabels, RDT_COLUMN_MAPPING, RDT_PHASE_COLUMN, 
+    RDT_IMAG_COLUMN, RDT_REAL_COLUMN, RDT_AMPLITUDE_COLUMN
+)
 from omc3.optics_measurements.constants import EXT
-from omc3.plotting.plot_optics_measurements import (_get_x_axis_column_and_label, _get_ip_positions,
-                                                    get_optics_style_params, get_plottfs_style_params)
+from omc3.plotting.plot_optics_measurements import (
+    _get_x_axis_column_and_label, _get_ip_positions,
+    get_optics_style_params, get_plottfs_style_params
+)
 from omc3.plotting.plot_tfs import plot as plot_tfs, get_full_output_path
-from omc3.plotting.utils import (annotations as pannot)
+from omc3.plotting.utils.windows import (
+    PlotWidget, TabWidget, VerticalTabWindow, 
+    log_no_qtpy_many_windows, create_pyplot_window_from_fig
+)
+from omc3.plotting.utils import annotations as pannot
 from omc3.utils import logging_tools
-from omc3.utils.iotools import PathOrStr
+from omc3.utils.iotools import PathOrStr, save_config
 
 LOG = logging_tools.get_logger(__name__)
 
 SPLIT_ID = "#_#"  # will appear in the figure ID, but should be fine to read
 PREFIX = "plot_corrections"
 
+
 def get_plotting_params() -> EntryPointParameters:
     params = EntryPointParameters()
     params.add_parameter(name="input_dir",
@@ -220,6 +234,12 @@ def get_plotting_style_parameters():
 def plot_checked_corrections(opt: DotDict):
     """ Entrypoint for the plotting function. """
     LOG.info("Plotting checked corrections.")
+
+    opt.input_dir = Path(opt.input_dir)
+    if opt.output_dir:
+        opt.output_dir = Path(opt.output_dir)
+        save_config(opt.output_dir, opt, __file__)
+
     # Preparations -------------------------------------------------------------
     correction_dirs: Dict[str, Path] = {}
     if len(opt.corrections) == 1 and not opt.corrections[0]:
@@ -248,7 +268,7 @@ def plot_checked_corrections(opt: DotDict):
             # get F1### column map without the I, R, A, P part based on the rdt-filename:
             LOG.debug(f"Plotting coupling correction for {filename}")
             new_figs = {}
-            for idx, y_colmap in enumerate(RDT_COLUMN_MAPPING.values()):  # AMP, PHASE, REAL or IMAG as column-map
+            for y_colmap in RDT_COLUMN_MAPPING.values():  # AMP, PHASE, REAL or IMAG as column-map
                 y_colmap = ColumnsAndLabels(
                     _column=y_colmap.column,
                     _label=y_colmap.label.format(filename),  # this one needs additional info
@@ -281,11 +301,27 @@ def plot_checked_corrections(opt: DotDict):
         fig_dict.update(new_figs)
 
     # Output -------------------------------------------------------------------
-    save_plots(opt.output_dir, figure_dict=fig_dict, input_dir=opt.input_dir if opt.individual_to_input else None)
-    show_plots(opt.show)
+    if opt.output_dir:
+        save_plots(
+            opt.output_dir, 
+            figure_dict=fig_dict, 
+            input_dir=opt.input_dir if opt.individual_to_input else None
+        )
+
+    if opt.show:
+        show_plots(fig_dict)
     return fig_dict
 
-def _create_correction_plots_per_filename(filename, measurements, correction_dirs, x_colmap, y_colmap, ip_positions, opt):
+
+def _create_correction_plots_per_filename(
+        filename: str, 
+        measurements: Path, 
+        correction_dirs: Dict[str, Path], 
+        x_colmap: ColumnsAndLabels, 
+        y_colmap: ColumnsAndLabels, 
+        ip_positions: Union[str, Dict[str, float], Path], 
+        opt: DotDict
+    ):
     """ Plot measurements and all different correction scenarios into a single plot. """
     full_filename = f"{filename}{EXT}"
     file_label = filename
@@ -305,7 +341,7 @@ def _create_correction_plots_per_filename(filename, measurements, correction_dir
         x_labels=[x_colmap.label],
         vertical_lines=ip_positions + opt.lines_manual,
         same_axes=["files"],
-        output_prefix=f"{PREFIX}_{file_label}_",  # used in the id, which is the fig_dict key
+        output_prefix=f"{file_label}_",  # used in the id, which is the fig_dict key
         **opt.get_subdict([
             'plot_styles', 'manual_style',
             'change_marker', 'errorbar_alpha',
@@ -366,41 +402,110 @@ def _create_correction_plots_per_filename(filename, measurements, correction_dir
     return figs
 
 
-def show_plots(show: bool):
-    """ Show plots if so desired. """
-    # plt.show()
-    if show:
-        plt.show()
-
-
-def save_plots(output_dir, figure_dict, input_dir=None):
+def save_plots(output_dir: Path, figure_dict: Dict[str, Figure], input_dir: Path = None):
     """ Save the plots. """
-    if not output_dir and not input_dir:
-        return
-
     for figname, fig in figure_dict.items():
         outdir = output_dir
         figname_parts = figname.split(SPLIT_ID)
         if len(figname_parts) == 1:  # no SPLIT_ID
             # these are the combined plots. They have the column name at the end,
             # which we do not care for here at the moment.
             # In case of multiple columns per file, this could be brought back
-            # (then we would also not need the RDT check).
-            figname = "_".join(figname.split("_")[:-1])
+            figname = "_".join([PREFIX] + figname.split("_")[:-1])
         else:
             # this is then the individual plots
             if input_dir:
                 # files go directly into the correction-scenario folders
                 outdir = input_dir / figname_parts[0]
                 figname = f"{PREFIX}_{figname_parts[1]}"
             else:
-                # everything goes into the output-dir (if given), but needs prefix
+                # everything goes into the output-dir (if given),
+                # but with correction-name as additional prefix
                 figname = "_".join([PREFIX] + figname_parts)
 
         output_path = get_full_output_path(outdir, figname)
-        if output_path is not None:
-            LOG.info(f"Saving Corrections Plot to '{output_path}'")
-            fig.savefig(output_path)
+        LOG.debug(f"Saving corrections plot to '{output_path}'")
+        fig.savefig(output_path)
+    
+    if input_dir:
+        LOG.info(f"Saved all correction plots in '{output_dir}'\n"
+                 f"and into the correction-scenario in '{input_dir}'.")
+    else:
+        LOG.info(f"Saved all correction plots in '{output_dir}'.")
+
+
+def show_plots(figure_dict: Dict[str, Figure]):
+    """Displays the provided figures.
+    If `qtpy` is installed, they are shown in a single window.
+    The individual corrections are sorted into vertical tabs,
+    the optics parameter into horizontal tabs. 
+    If `qtpy` is not installed, they are simply shown as individual figures.
+    This is not recommended
+    """
+    try:
+        window = VerticalTabWindow("Correction Check")
+    except TypeError:
+        log_no_qtpy_many_windows()
+        for name, fig in figure_dict.items():
+            create_pyplot_window_from_fig(fig, title=name.replace(SPLIT_ID, " "))
+        plt.show()
+        return
+
+    rdt_pattern = re.compile(r"f\d{4}")
+    rdt_complement = {
+        RDT_REAL_COLUMN.text_label: RDT_IMAG_COLUMN,
+        RDT_AMPLITUDE_COLUMN.text_label: RDT_PHASE_COLUMN,
+    }
+
+    correction_names = sorted(set([k.split(SPLIT_ID)[0] for k in figure_dict.keys() if SPLIT_ID in k]))
+    for correction_name in [None] + list(correction_names):
+        if not correction_name:
+            parameter_names = iter(sorted(k for k in figure_dict.keys() if SPLIT_ID not in k))
+            correction_tab_name = "All Corrections"
+        else:
+            parameter_names = iter(sorted(k for k in figure_dict.keys() if correction_name in k))
+            correction_tab_name = correction_name
+
+        current_tab = TabWidget(title=correction_tab_name)
+        window.add_tab(current_tab)
+
+        for name_x in parameter_names:
+            # extract the filename (and column-name in case of multi-correction-file)
+            tab_prename = name_x.split(SPLIT_ID)[-1] 
+
+            if rdt_pattern.match(tab_prename):
+                # Handle RDTs: Get the rdt column and if it's amplitude or real, 
+                # we look for the respective complement column (phase, imag).
+                # Both, column and complement column are then added to the tab,
+                # which is named after the rdt followed by either AP (amp/phase) or RI (real/imag)).
+                rdt, column = tab_prename.split("_")[:2]
+                try:
+                    complement_column: ColumnsAndLabels = rdt_complement[column]
+                except KeyError:
+                    # skip phase and imag as they will become name_y for amp and real.
+                    continue
+
+                if not correction_name:
+                    name_y = "_".join([rdt, complement_column.text_label, complement_column.expected_column])
+                else:
+                    name_y = "_".join(name_x.split("_")[:-1] + [complement_column.text_label,])
+
+                tab_name = f"{rdt} {column[0].upper()}/{complement_column.text_label[0].upper()}"
+
+            else:
+                # Handle non-RDT columns: As they are sorted alphabetically, the current column
+                # is x and the following column is y. They are added to the tab, which 
+                # is named by the optics parameter without plane.
+                tab_name = " ".join(tab_prename.split("_")[:-1 if correction_name else -2])  # remove plane (and column-name)
+                name_y = next(parameter_names)
+
+            new_tab = PlotWidget(
+                figure_dict[name_x],
+                figure_dict[name_y],
+                title=tab_name,
+            )
+            current_tab.add_tab(new_tab)
+    window.show()
 
 
 def _get_corrected_measurement_names(correction_dirs: Iterable[Path]) -> Set[str]:
@@ -415,3 +520,7 @@ def _get_corrected_measurement_names(correction_dirs: Iterable[Path]) -> Set[str
         tfs_files &= new_files
     # tfs_files -= {Path(MODEL_MATCHED_FILENAME).stem, Path(MODEL_NOMINAL_FILENAME).stem}  # no need, filtered later anyway
     return tfs_files
+
+
+if __name__ == "__main__":
+    plot_checked_corrections()

omc3/plotting/plot_tfs.py

@@ -105,26 +105,31 @@
 - **y_lim** *(float, int None)*: Limits on the y axis (Tupel)
 """
 from collections import OrderedDict
-from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Optional, Dict
+from typing import Dict
 
 import matplotlib
+from matplotlib import pyplot as plt, rcParams
 from matplotlib.axes import Axes
-from numpy.typing import ArrayLike
 
 import tfs
-from generic_parser import EntryPointParameters, entrypoint, DotDict
-from generic_parser.entry_datatypes import DictAsString, get_multi_class
-from matplotlib import pyplot as plt, rcParams
+from generic_parser import EntryPointParameters, entrypoint
+from generic_parser.entry_datatypes import DictAsString
 
 from omc3.definitions.constants import PLANES
 from omc3.optics_measurements.constants import EXT
 from omc3.plotting.optics_measurements.constants import DEFAULTS
 from omc3.plotting.optics_measurements.utils import FigureCollector, DataSet, IDMap, safe_format
+from omc3.plotting.utils.windows import (
+    PlotWidget, SimpleTabWindow, is_qtpy_installed, log_no_qtpy_many_windows, create_pyplot_window_from_fig
+)
 from omc3.plotting.spectrum.utils import get_unique_filenames, output_plot
-from omc3.plotting.utils import (annotations as pannot, lines as plines,
-                                 style as pstyle, colors as pcolors)
+from omc3.plotting.utils import (
+    annotations as pannot, 
+    lines as plines,
+    style as pstyle, 
+    colors as pcolors,
+)
 from omc3.plotting.utils.lines import VERTICAL_LINES_TEXT_LOCATIONS
 from omc3.utils.iotools import PathOrStr, save_config, OptionalStr, OptionalFloat
 from omc3.utils.logging_tools import get_logger, list2str
@@ -302,7 +307,7 @@ def get_params():
 @entrypoint(get_params(), strict=True)
 def plot(opt):
     """Main plotting function."""
-    LOG.info(f"Starting plotting of tfs files: {list2str(opt.files):s}")
+    LOG.debug(f"Starting plotting of tfs files: {list2str(opt.files):s}")
     if opt.output is not None:
         save_config(Path(opt.output), opt, __file__)
 
@@ -531,8 +536,22 @@ def _create_plots(fig_collection, opt):
         output_plot(fig_container)
 
     if opt.show:
+        if len(fig_collection) > 1 and is_qtpy_installed():
+            window = SimpleTabWindow("Tfs Plots")
+            for fig_container in fig_collection.figs.values():
+                tab = PlotWidget(fig_container.fig, title=fig_container.id)
+                window.add_tab(tab)
+            window.show()
+            return
+        
+        if len(fig_collection) > rcParams['figure.max_open_warning']:
+            log_no_qtpy_many_windows()
+
+        for fig_container in fig_collection.figs.values():
+            create_pyplot_window_from_fig(fig_container.fig, title=fig_container.id)
         plt.show()
 
+            
 
 def _plot_data(ax: Axes, data: Dict[str, DataSet], change_marker: bool, ebar_alpha: float):
     for idx, (label, values) in enumerate(data.items()):