spectre-core 0.0.22__py3-none-any.whl → 0.0.23__py3-none-any.whl

spectre_core/plotting/_panel_stack.py

@@ -15,13 +15,10 @@ from spectre_core.spectrograms import TimeType
 from spectre_core.config import TimeFormat, get_batches_dir_path
 from ._base import BasePanel, XAxisType
 from ._format import PanelFormat
-from ._panels import (
-    PanelName, SpectrogramPanel, TimeCutsPanel, FrequencyCutsPanel
-)
+from ._panels import PanelName, SpectrogramPanel, TimeCutsPanel, FrequencyCutsPanel
 
-def _is_cuts_panel(
-    panel: BasePanel
-) -> bool:
+
+def _is_cuts_panel(panel: BasePanel) -> bool:
     """Check if a panel contains spectrogram cuts.
 
     :param panel: The panel to check.
@@ -30,9 +27,7 @@ def _is_cuts_panel(
     return panel.name in {PanelName.FREQUENCY_CUTS, PanelName.TIME_CUTS}
 
 
-def _is_spectrogram_panel(
-    panel: BasePanel
-) -> bool:
+def _is_spectrogram_panel(panel: BasePanel) -> bool:
     """Check if a panel is a spectrogram panel.
 
     :param panel: The panel to check.
@@ -43,12 +38,13 @@ def _is_spectrogram_panel(
 
 class PanelStack:
     """Visualise spectrogram data in a stack of panels."""
+
     def __init__(
-        self, 
+        self,
         panel_format: PanelFormat = PanelFormat(),
         time_type: TimeType = TimeType.RELATIVE,
         figsize: Tuple[int, int] = (15, 8),
-        non_interactive: bool = False
+        non_interactive: bool = False,
     ) -> None:
         """Initialize an instance of `PanelStack`.
 
@@ -59,92 +55,79 @@ class PanelStack:
         self._panel_format = panel_format
         self._time_type = time_type
         self._figsize = figsize
-        
+
         if non_interactive:
             # Use a non-interactive matplotlib backend, which can only write files.
-            use('agg')
+            use("agg")
 
-        self._panels             : list[BasePanel] = []
+        self._panels: list[BasePanel] = []
         self._superimposed_panels: list[BasePanel] = []
-        
+
         self._fig: Optional[Figure] = None
         self._axs: Optional[np.ndarray] = None
 
-
     @property
-    def time_type(
-        self
-    ) -> TimeType:
+    def time_type(self) -> TimeType:
         """The type of time we assign to the spectrograms"""
         return self._time_type
 
-
     @property
-    def panels(
-        self
-    ) -> list[BasePanel]:
+    def panels(self) -> list[BasePanel]:
         """Get the panels in the stack, sorted by their `XAxisType`."""
         return list(sorted(self._panels, key=lambda panel: panel.xaxis_type.value))
 
-
     @property
-    def fig(
-        self
-    ) -> Figure:
+    def fig(self) -> Figure:
         """Get the shared `matplotlib` figure for the panel stack.
 
         :raises ValueError: If the axes have not been initialized.
         """
         if self._fig is None:
-            raise ValueError(f"An unexpected error has occured, `fig` must be set for the panel stack.")
+            raise ValueError(
+                f"An unexpected error has occured, `fig` must be set for the panel stack."
+            )
         return self._fig
 
-
     @property
-    def axs(
-        self
-    ) -> np.ndarray:
+    def axs(self) -> np.ndarray:
         """Get the `matplotlib` axes in the stack.
 
         :return: An array of `matplotlib.axes.Axes`, one for each panel in the stack.
         :raises ValueError: If the axes have not been initialized.
         """
         if self._axs is None:
-            raise ValueError(f"An unexpected error has occured, `axs` must be set for the panel stack.")
+            raise ValueError(
+                f"An unexpected error has occured, `axs` must be set for the panel stack."
+            )
         return np.atleast_1d(self._axs)
 
-
     @property
-    def num_panels(
-        self
-    ) -> int:
+    def num_panels(self) -> int:
         """Get the number of panels in the stack."""
         return len(self._panels)
 
-
     def add_panel(
-        self, 
+        self,
         panel: BasePanel,
         identifier: Optional[str] = None,
-        panel_format: Optional[PanelFormat] = None
+        panel_format: Optional[PanelFormat] = None,
     ) -> None:
         """Add a panel to the stack.
 
         :param panel: An instance of a `BasePanel` subclass to be added to the stack.
         :param identifier: An optional string to link the panel with others for superimposing.
         """
-        panel.panel_format = panel_format or self._panel_format 
-        panel.time_type    = self._time_type
-        if identifier: 
+        panel.panel_format = panel_format or self._panel_format
+        panel.time_type = self._time_type
+        if identifier:
             panel.identifier = identifier
         self._panels.append(panel)
 
-
     def superimpose_panel(
-        self, 
+        self,
         panel: BasePanel,
         identifier: Optional[str] = None,
-        panel_format: Optional[PanelFormat] = None
+        panel_format: Optional[PanelFormat] = None,
     ) -> None:
         """Superimpose a panel onto an existing panel in the stack.
 
@@ -154,49 +137,42 @@ class PanelStack:
         if identifier:
             panel.identifier = identifier
         panel.panel_format = panel_format or self._panel_format
-        panel.time_type    = self._time_type
+        panel.time_type = self._time_type
         self._superimposed_panels.append(panel)
 
-
-    def _init_plot_style(
-        self
-    ) -> None:
+    def _init_plot_style(self) -> None:
         """Initialize the global plot style for the stack.
 
         This method sets `matplotlib` styles and font sizes based on the `panel_format`.
         """
         plt.style.use(self._panel_format.style)
-        
-        plt.rc('font', size=self._panel_format.small_size)
-        
-        plt.rc('axes', titlesize=self._panel_format.medium_size, 
-                       labelsize=self._panel_format.medium_size)
-        
-        plt.rc('xtick', labelsize=self._panel_format.small_size)
-        plt.rc('ytick', labelsize=self._panel_format.small_size)
-        
-        plt.rc('legend', fontsize=self._panel_format.small_size)
-        plt.rc('figure', titlesize=self._panel_format.large_size)
-
-
-    def _create_figure_and_axes(
-        self
-    ) -> None:
+
+        plt.rc("font", size=self._panel_format.small_size)
+
+        plt.rc(
+            "axes",
+            titlesize=self._panel_format.medium_size,
+            labelsize=self._panel_format.medium_size,
+        )
+
+        plt.rc("xtick", labelsize=self._panel_format.small_size)
+        plt.rc("ytick", labelsize=self._panel_format.small_size)
+
+        plt.rc("legend", fontsize=self._panel_format.small_size)
+        plt.rc("figure", titlesize=self._panel_format.large_size)
+
+    def _create_figure_and_axes(self) -> None:
         """Create the `matplotlib` figure and axes for the panel stack.
 
         This initializes a figure with a specified number of vertically stacked axes.
         """
-        self._fig, self._axs = plt.subplots(self.num_panels, 
-                                            1, 
-                                            figsize=self._figsize, 
-                                            layout="constrained")
-
+        self._fig, self._axs = plt.subplots(
+            self.num_panels, 1, figsize=self._figsize, layout="constrained"
+        )
 
-    def _assign_axes(
-        self
-    ) -> None:
+    def _assign_axes(self) -> None:
         """Assign each axes in the figure to some panel in the stack.
-        
+
         Axes are shared between panels with common `xaxis_type`.
         """
         shared_axes: dict[XAxisType, Axes] = {}
@@ -204,15 +180,11 @@ class PanelStack:
             panel.ax = self.axs[i]
             panel.fig = self._fig
             if panel.xaxis_type in shared_axes:
-                panel.ax.sharex( shared_axes[panel.xaxis_type] )
+                panel.ax.sharex(shared_axes[panel.xaxis_type])
             else:
                 shared_axes[panel.xaxis_type] = panel.ax
 
-
-    def _overlay_cuts(
-        self, 
-        cuts_panel: BasePanel
-    ) -> None:
+    def _overlay_cuts(self, cuts_panel: BasePanel) -> None:
         """Overlay cuts onto corresponding spectrogram panels.
 
         For a given cuts panel, locate any spectrogram panels in the stack with a matching tag
@@ -224,82 +196,77 @@ class PanelStack:
             if _is_spectrogram_panel(panel) and panel.tag == cuts_panel.tag:
                 panel = cast(SpectrogramPanel, panel)
                 if cuts_panel.name == PanelName.FREQUENCY_CUTS:
-                    panel.overlay_frequency_cuts( cast(FrequencyCutsPanel, cuts_panel) )
+                    panel.overlay_frequency_cuts(cast(FrequencyCutsPanel, cuts_panel))
                 elif cuts_panel.name == PanelName.TIME_CUTS:
-                    panel.overlay_time_cuts( cast(TimeCutsPanel, cuts_panel) )
-
+                    panel.overlay_time_cuts(cast(TimeCutsPanel, cuts_panel))
 
-    def _overlay_superimposed_panels(
-        self
-    ) -> None:
+    def _overlay_superimposed_panels(self) -> None:
         """Superimpose panels onto matching panels in the stack.
 
-        For each superimposed panel, find a matching panel in the stack with the same name 
-        and identifier. Share the axes and figure, then draw the superimposed panel. If the 
+        For each superimposed panel, find a matching panel in the stack with the same name
+        and identifier. Share the axes and figure, then draw the superimposed panel. If the
         panel contains cuts, overlay those cuts onto the corresponding spectrogram panels.
         """
         for super_panel in self._superimposed_panels:
             for panel in self._panels:
-                if panel.name == super_panel.name and (panel.identifier == super_panel.identifier):
+                if panel.name == super_panel.name and (
+                    panel.identifier == super_panel.identifier
+                ):
                     super_panel.ax, super_panel.fig = panel.ax, self._fig
                     super_panel.draw()
                     if _is_cuts_panel(super_panel):
                         self._overlay_cuts(super_panel)
 
-    def _make_figure(
-        self
-    ) -> None:
+    def _make_figure(self) -> None:
         """Make the panel stack figure."""
         if self.num_panels < 1:
             raise ValueError(f"There must be at least one panel in the stack.")
-        
+
         self._init_plot_style()
         self._create_figure_and_axes()
         self._assign_axes()
-        
+
         last_panel_per_axis = {panel.xaxis_type: panel for panel in self.panels}
-        
+
         for panel in self.panels:
             panel.draw()
             panel.annotate_yaxis()
-            
+
             if panel == last_panel_per_axis[panel.xaxis_type]:
                 panel.annotate_xaxis()
             else:
                 panel.hide_xaxis_labels()
-                
+
             if _is_cuts_panel(panel):
                 self._overlay_cuts(panel)
-                
+
         self._overlay_superimposed_panels()
 
-    def show(
-        self
-    ) -> None:
+    def show(self) -> None:
         """Display the panel stack figure."""
         self._make_figure()
         plt.show()
-        
-        
+
     def save(
         self,
     ) -> str:
-        """Save the panel stack figure as a batch file under the tag of the first 
-        panel in the stack. The file format is `png`. 
-        
+        """Save the panel stack figure as a batch file under the tag of the first
+        panel in the stack. The file format is `png`.
+
         :return: The file path of the newly created batch file containing the figure.
         """
         self._make_figure()
         first_panel = self._panels[0]
-        
-        start_dt = cast(datetime, first_panel.spectrogram.start_datetime.astype(datetime))
-        batch_name = f"{start_dt.strftime(TimeFormat.DATETIME)}_{first_panel.spectrogram.tag}"
+
+        start_dt = cast(
+            datetime, first_panel.spectrogram.start_datetime.astype(datetime)
+        )
+        batch_name = (
+            f"{start_dt.strftime(TimeFormat.DATETIME)}_{first_panel.spectrogram.tag}"
+        )
         batch_file_path = os.path.join(
             get_batches_dir_path(start_dt.year, start_dt.month, start_dt.day),
-            f"{batch_name}.png"
+            f"{batch_name}.png",
         )
         plt.savefig(batch_file_path)
         return batch_file_path
-        
-        
-