spectre-core 0.0.9__py3-none-any.whl → 0.0.10__py3-none-any.whl

spectre_core/plotting/{panel_stack.py → _panel_stack.py}

@@ -3,25 +3,36 @@
 # SPDX-License-Identifier: GPL-3.0-or-later
 
 import numpy as np
-from typing import List, Optional, Tuple
+from typing import Optional, Tuple
 import matplotlib.pyplot as plt
 from matplotlib.figure import Figure
 from matplotlib.axes import Axes
 
-from spectre_core.spectrograms.spectrogram import Spectrogram
-from spectre_core.plotting.base import BasePanel, CutsPanel
-from spectre_core.plotting.format import DEFAULT_FORMATS
-from spectre_core.plotting.factory import get_panel
-from spectre_core.plotting.library.spectrogram.panel import Panel as SpectrogramPanel
+from spectre_core.spectrograms import TimeTypes
+from ._base import BasePanel
+from ._format import PanelFormat, DEFAULT_FORMAT
+from ._panels import PanelNames
+
+
+def _is_cuts_panel(panel: BasePanel) -> bool:
+    return (panel.name == PanelNames.FREQUENCY_CUTS or panel.name == PanelNames.TIME_CUTS)
+
+
+def _is_spectrogram_panel(panel: BasePanel) -> bool:
+    return (panel.name == PanelNames.SPECTROGRAM)
+
 
 class PanelStack:
     def __init__(self, 
-                 time_type: str = "seconds", 
+                 panel_format: PanelFormat = DEFAULT_FORMAT,
+                 time_type: str = TimeTypes.SECONDS,
                  figsize: Tuple[int, int] = (10, 10)):
+        self._panel_format = panel_format
         self._time_type = time_type
         self._figsize = figsize
-        self._panels: List[BasePanel] = []
-        self._superimposed_panels: List[BasePanel] = []
+
+        self._panels: list[BasePanel] = []
+        self._superimposed_panels: list[BasePanel] = []
         self._fig: Optional[Figure] = None
         self._axs: Optional[np.ndarray[Axes]] = None
 
@@ -32,7 +43,7 @@ class PanelStack:
 
 
     @property
-    def panels(self) -> List[BasePanel]:
+    def panels(self) -> list[BasePanel]:
         return sorted(self._panels, key=lambda panel: panel.x_axis_type)
 
 
@@ -52,50 +63,41 @@ class PanelStack:
 
 
     def add_panel(self, 
-                  panel_name: str, 
-                  spectrogram: Spectrogram,
-                  *args, 
-                  identifier: Optional[str] = None,
-                  **kwargs) -> None:
-        panel = get_panel(panel_name, 
-                          spectrogram, 
-                          self._time_type, 
-                          *args, 
-                          **kwargs)
+                  panel: BasePanel,
+                  identifier: Optional[str] = None) -> None:
+        panel.panel_format = self._panel_format
+        panel.time_type    = self._time_type
         if identifier: 
             panel.identifier = identifier
         self._panels.append(panel)
 
 
     def superimpose_panel(self, 
-                          panel_name: str, 
-                          spectrogram: Spectrogram, 
-                          *args, 
-                          identifier: Optional[str] = None,
-                          **kwargs) -> None:
-        panel = get_panel(panel_name, 
-                          spectrogram, 
-                          self._time_type, 
-                          *args, 
-                          **kwargs)
+                          panel: BasePanel,
+                          identifier: Optional[str] = None) -> None:
         if identifier:
             panel.identifier = identifier
+        panel.panel_format = self._panel_format
         self._superimposed_panels.append(panel)
 
 
     def _init_plot_style(self) -> None:
-        plt.style.use(DEFAULT_FORMATS.style)
-        plt.rc('font', size=DEFAULT_FORMATS.small_size)
-        plt.rc('axes', titlesize=DEFAULT_FORMATS.medium_size, 
-                       labelsize=DEFAULT_FORMATS.medium_size)
-        plt.rc('xtick', labelsize=DEFAULT_FORMATS.small_size)
-        plt.rc('ytick', labelsize=DEFAULT_FORMATS.small_size)
-        plt.rc('legend', fontsize=DEFAULT_FORMATS.small_size)
-        plt.rc('figure', titlesize=DEFAULT_FORMATS.large_size)
+        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:
-        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:
@@ -109,10 +111,10 @@ class PanelStack:
                 shared_axes[panel.x_axis_type] = panel.ax
 
 
-    def _overlay_cuts(self, cuts_panel: CutsPanel) -> None:
+    def _overlay_cuts(self, cuts_panel: BasePanel) -> None:
         """Given a cuts panel, finds any corresponding spectrogram panels and adds the appropriate overlay"""
         for panel in self.panels:
-            is_corresponding_panel = isinstance(panel, SpectrogramPanel) and (panel.tag == cuts_panel.tag)
+            is_corresponding_panel = _is_spectrogram_panel(panel) and (panel.tag == cuts_panel.tag)
             if is_corresponding_panel:
                 panel.overlay_cuts(cuts_panel)
 
@@ -120,10 +122,10 @@ class PanelStack:
     def _overlay_superimposed_panels(self) -> None:
         for super_panel in self._superimposed_panels:
             for panel in self._panels:
-                if panel.name == super_panel.name and (panel.identifier == super_panel.identifier): # find the matching panels via panel names
-                    super_panel.ax, super_panel.fig = panel.ax, self._fig # and superimpose via axes sharing
+                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 isinstance(super_panel, CutsPanel):
+                    if _is_cuts_panel(super_panel):
                         self._overlay_cuts(super_panel)
 
 
@@ -139,9 +141,7 @@ class PanelStack:
                 panel.annotate_x_axis()
             else:
                 panel.hide_x_axis_labels()
-            if isinstance(panel, CutsPanel):
-                        self._overlay_cuts(panel)
-
-
+            if _is_cuts_panel(panel):
+                self._overlay_cuts(panel)
         self._overlay_superimposed_panels()
         plt.show()

spectre_core/plotting/_panels.py

@@ -0,0 +1,234 @@
+# SPDX-FileCopyrightText: © 2024 Jimmy Fitzpatrick <jcfitzpatrick12@gmail.com>
+# This file is part of SPECTRE
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+from typing import Optional
+from datetime import datetime
+from dataclasses import dataclass
+from warnings import warn
+
+from matplotlib.colors import LogNorm
+import numpy as np
+
+from spectre_core.spectrograms import Spectrogram, FrequencyCut, TimeCut
+from ._base import BasePanel, BaseSpectrumPanel, BaseTimeSeriesPanel
+
+
+@dataclass(frozen=True)
+class PanelNames:
+    SPECTROGRAM            : str = "spectrogram"
+    FREQUENCY_CUTS         : str = "frequency_cuts"
+    TIME_CUTS              : str = "time_cuts"
+    INTEGRAL_OVER_FREQUENCY: str = "integral_over_frequency"
+
+
+class _FrequencyCutsPanel(BaseSpectrumPanel):
+    def __init__(self, 
+                 spectrogram: Spectrogram, 
+                 *times: list[float | str],
+                 dBb: bool = False,
+                 peak_normalise: bool = False):
+        super().__init__(PanelNames.FREQUENCY_CUTS,
+                         spectrogram)
+        self._times = times
+        self._dBb = dBb
+        self._peak_normalise = peak_normalise
+        # map each time cut to the corresponding FrequencyCut dataclass
+        self._frequency_cuts: Optional[dict[float | datetime, FrequencyCut]] = {}
+
+
+    @property
+    def frequency_cuts(self) -> dict[float | str, FrequencyCut]:
+        if not self._frequency_cuts:
+            for time in self._times:
+                frequency_cut = self._spectrogram.get_frequency_cut(time,
+                                                                    dBb = self._dBb,
+                                                                    peak_normalise = self._peak_normalise)
+                self._frequency_cuts[frequency_cut.time] = frequency_cut
+        return self._frequency_cuts
+
+
+    @property
+    def times(self) -> list[float | datetime]:
+        return list(self.frequency_cuts.keys())
+    
+
+    def draw(self):
+        for time, color in self.bind_to_colors():
+            frequency_cut = self.frequency_cuts[time]
+            self.ax.step(self.frequencies*1e-6, # convert to MHz
+                         frequency_cut.cut, 
+                         where='mid', 
+                         color = color)
+    
+
+    def annotate_y_axis(self) -> None:
+        if self._dBb:
+            self.ax.set_ylabel('dBb')
+        elif self._peak_normalise:
+            return # no y-axis label
+        else:
+            self.ax.set_ylabel(f'{self._spectrogram.spectrum_type.capitalize()}')
+
+    
+    def bind_to_colors(self):
+        return super().bind_to_colors(self.times, cmap = self.panel_format.cuts_cmap)
+    
+
+class _IntegralOverFrequencyPanel(BaseTimeSeriesPanel):
+    def __init__(self, 
+                 spectrogram: Spectrogram, 
+                 peak_normalise: bool = False,
+                 background_subtract: bool = False):
+        super().__init__(PanelNames.INTEGRAL_OVER_FREQUENCY,
+                         spectrogram)
+        self._peak_normalise = peak_normalise
+        self._background_subtract = background_subtract
+
+
+    def draw(self):
+        I = self._spectrogram.integrate_over_frequency(correct_background = self._background_subtract,
+                                                       peak_normalise = self._peak_normalise)
+        self.ax.step(self.times, I, where="mid", color = self.panel_format.integral_color)
+ 
+
+    def annotate_y_axis(self):
+        # no y-axis label
+        return 
+
+
+class _TimeCutsPanel(BaseTimeSeriesPanel):
+    def __init__(self, 
+                 spectrogram: Spectrogram, 
+                 *frequencies: list[float],
+                 dBb: bool = False,
+                 peak_normalise: bool = False,
+                 background_subtract: bool = False):
+        super().__init__(PanelNames.TIME_CUTS, 
+                         spectrogram)
+        self._frequencies = frequencies
+        self._dBb = dBb
+        self._peak_normalise = peak_normalise
+        self._background_subtract = background_subtract
+        # map each cut frequency to the corresponding TimeCut dataclass
+        self._time_cuts: Optional[dict[float, TimeCut]] = {} 
+    
+
+    @property
+    def time_cuts(self) -> dict[float, TimeCut]:
+        if not self._time_cuts:
+            for frequency in self._frequencies:
+                time_cut =  self._spectrogram.get_time_cut(frequency,
+                                                           dBb = self._dBb,
+                                                           peak_normalise = self._peak_normalise,
+                                                           correct_background = self._background_subtract,
+                                                           return_time_type=self._time_type)
+                self._time_cuts[time_cut.frequency] = time_cut
+        return self._time_cuts
+    
+
+    @property
+    def frequencies(self) -> list[float]:
+        return list(self.time_cuts.keys())
+
+
+    def draw(self):
+        for frequency, color in self.bind_to_colors():
+            time_cut = self.time_cuts[frequency]
+            self.ax.step(self.times, 
+                         time_cut.cut, 
+                         where='mid', 
+                         color = color)
+    
+
+    def annotate_y_axis(self) -> None:
+        if self._dBb:
+            self.ax.set_ylabel('dBb')
+        elif self._peak_normalise:
+            return # no y-axis label
+        else:
+            self.ax.set_ylabel(f'{self._spectrogram.spectrum_type.capitalize()}')
+
+    
+    def bind_to_colors(self):
+        return super().bind_to_colors(self.frequencies, cmap = self.panel_format.cuts_cmap)
+    
+
+class _SpectrogramPanel(BaseTimeSeriesPanel):
+    def __init__(self, 
+                 spectrogram: Spectrogram, 
+                 log_norm: bool = False,
+                 dBb: bool = False,
+                 vmin: float | None = -1,
+                 vmax: float | None = 2):
+        super().__init__(PanelNames.SPECTROGRAM,
+                         spectrogram)
+        self._log_norm = log_norm
+        self._dBb = dBb
+        self._vmin = vmin
+        self._vmax = vmax
+
+
+    def draw(self):
+        dynamic_spectra = self._spectrogram.dynamic_spectra_as_dBb if self._dBb else self._spectrogram.dynamic_spectra
+
+        norm = LogNorm(vmin=np.nanmin(dynamic_spectra[dynamic_spectra > 0]), 
+                       vmax=np.nanmax(dynamic_spectra)) if self._log_norm else None
+        
+
+        if self._log_norm and (self._vmin or self._vmax):
+            warn("vmin/vmax will be ignored while using log_norm.")
+            self._vmin = None 
+            self._vmax = None
+        
+        # Plot the spectrogram with kwargs
+        pcm = self.ax.pcolormesh(self.times, 
+                            self._spectrogram.frequencies * 1e-6, 
+                            dynamic_spectra,
+                            vmin=self._vmin, 
+                            vmax=self._vmax,
+                            norm=norm, 
+                            cmap=self.panel_format.spectrogram_cmap)
+        
+        # Add colorbar if dBb is used
+        if self._dBb:
+            cbar = self.fig.colorbar(pcm, 
+                                      ax=self.ax, 
+                                      ticks=np.linspace(self._vmin, self._vmax, 6, dtype=int))
+            cbar.set_label('dBb')
+
+
+    def annotate_y_axis(self) -> None:
+        self.ax.set_ylabel('Frequency [MHz]')
+        return
+    
+    
+    def overlay_cuts(self, cuts_panel: BasePanel) -> None:
+        if cuts_panel.name == PanelNames.TIME_CUTS:
+            self._overlay_time_cuts(cuts_panel)
+        elif cuts_panel.name == PanelNames.FREQUENCY_CUTS:
+            self._overlay_frequency_cuts(cuts_panel)
+
+
+    def _overlay_time_cuts(self, time_cuts_panel: _TimeCutsPanel) -> None:
+        for frequency, color in time_cuts_panel.bind_to_colors():
+            self.ax.axhline(frequency*1e-6, # convert to MHz
+                            color = color,
+                            linewidth=self.panel_format.line_width
+                            )
+            
+            
+    def _overlay_frequency_cuts(self, frequency_cuts_panel: _FrequencyCutsPanel) -> None:
+        for time, color in frequency_cuts_panel.bind_to_colors():
+            self.ax.axvline(time,
+                            color = color,
+                            linewidth=self.panel_format.line_width
+                            )
+
+
+@dataclass(frozen=True)
+class Panels:
+    TimeCuts                    = _TimeCutsPanel
+    FrequencyCuts               = _FrequencyCutsPanel
+    Spectrogram                 = _SpectrogramPanel
+    IntegralOverFrequency       = _IntegralOverFrequencyPanel

spectre_core/post_processing/__init__.py

@@ -2,5 +2,13 @@
 # This file is part of SPECTRE
 # SPDX-License-Identifier: GPL-3.0-or-later
 
-# dynamically import all event handlers
-import spectre_core.post_processing.library
+# event handler class decorators take effect on import
+from .library._fixed_center_frequency import _EventHandler
+from .library._swept_center_frequency import _EventHandler
+
+from ._factory import get_event_handler_from_tag
+from ._post_processor import PostProcessor
+
+__all__ = [
+    "PostProcessor", "get_event_handler_from_tag"
+]

spectre_core/post_processing/_base.py

@@ -0,0 +1,119 @@
+# SPDX-FileCopyrightText: © 2024 Jimmy Fitzpatrick <jcfitzpatrick12@gmail.com>
+# This file is part of SPECTRE
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+from logging import getLogger
+_LOGGER = getLogger(__name__)
+
+from typing import Optional
+from abc import ABC, abstractmethod
+from scipy.signal import ShortTimeFFT, get_window
+
+from watchdog.events import FileSystemEventHandler, FileCreatedEvent
+
+from spectre_core.capture_configs import CaptureConfig, PNames
+from spectre_core.chunks import BaseChunk, get_chunk_from_tag
+from spectre_core.spectrograms import Spectrogram, join_spectrograms
+
+
+def make_sft_instance(capture_config: CaptureConfig
+) -> ShortTimeFFT:
+    sample_rate   = capture_config.get_parameter_value(PNames.SAMPLE_RATE)
+    window_hop    = capture_config.get_parameter_value(PNames.WINDOW_HOP)
+    window_type   = capture_config.get_parameter_value(PNames.WINDOW_TYPE)
+    window_size   = capture_config.get_parameter_value(PNames.WINDOW_SIZE)
+    window = get_window(window_type, 
+                        window_size)
+    return ShortTimeFFT(window, 
+                        window_hop,
+                        sample_rate, 
+                        fft_mode = "centered")
+
+
+class BaseEventHandler(ABC, FileSystemEventHandler):
+    def __init__(self, 
+                 tag: str):
+        self._tag = tag
+
+        # the tag tells us 'what type' of data is stored in the files for each chunk
+        self._Chunk: BaseChunk = get_chunk_from_tag(tag)
+        # load the capture config corresponding to the tag
+        self._capture_config   = CaptureConfig(tag)
+
+        # post processing is triggered by files with this extension
+        self._watch_extension = self._capture_config.get_parameter_value(PNames.WATCH_EXTENSION)
+
+        # store the next file to be processed (specifically, the absolute file path of the file)
+        self._queued_file: Optional[str] = None
+
+        # store batched spectrograms as they are created into a cache
+        # which is flushed periodically according to a user defined 
+        # time range
+        self._cached_spectrogram: Optional[Spectrogram] = None
+
+
+    @abstractmethod
+    def process(self, 
+                absolute_file_path: str) -> None:
+        """Process the file stored at the input absolute file path.
+        
+        To be implemented by derived classes.
+        """
+
+    def on_created(self, 
+                   event: FileCreatedEvent):
+        """Process a newly created batch file, only once the next batch is created.
+        
+        Since we assume that the batches are non-overlapping in time, this guarantees
+        we avoid post processing a file while it is being written to. Files are processed
+        sequentially, in the order they are created.
+        """
+
+        # the 'src_path' attribute holds the absolute path of the newly created file
+        absolute_file_path = event.src_path
+        
+        # only 'notice' a file if it ends with the appropriate extension
+        # as defined in the capture config
+        if absolute_file_path.endswith(self._watch_extension):
+            _LOGGER.info(f"Noticed {absolute_file_path}")
+            
+            # If there exists a queued file, try and process it
+            if self._queued_file is not None:
+                try:
+                    self.process(self._queued_file)
+                except Exception:
+                    _LOGGER.error(f"An error has occured while processing {self._queued_file}",
+                                  exc_info=True)
+                     # flush any internally stored spectrogram on error to avoid lost data
+                    self._flush_cache()
+                    # re-raise the exception to the main thread
+                    raise
+            
+            # Queue the current file for processing next
+            _LOGGER.info(f"Queueing {absolute_file_path} for post processing")
+            self._queued_file = absolute_file_path
+    
+
+    def _cache_spectrogram(self, 
+                           spectrogram: Spectrogram) -> None:
+        _LOGGER.info("Joining spectrogram")
+
+        if self._cached_spectrogram is None:
+            self._cached_spectrogram = spectrogram
+        else:
+            self._cached_spectrogram = join_spectrograms([self._cached_spectrogram, spectrogram])
+
+        # if the time range is not specified
+        time_range = self._capture_config.get_parameter_value(PNames.TIME_RANGE) or self._capture_config.get_parameter_value(PNames.BATCH_SIZE)
+  
+        if self._cached_spectrogram.time_range >= time_range:
+            self._flush_cache()
+    
+
+    def _flush_cache(self) -> None:
+        if self._cached_spectrogram:
+            _LOGGER.info(f"Flushing spectrogram to file with chunk start time "
+                         f"'{self._cached_spectrogram.chunk_start_time}'")
+            self._cached_spectrogram.save()
+            _LOGGER.info("Flush successful, resetting spectrogram cache")
+            self._cached_spectrogram = None # reset the cache

spectre_core/post_processing/{factory.py → _factory.py}

@@ -2,12 +2,13 @@
 # This file is part of SPECTRE
 # SPDX-License-Identifier: GPL-3.0-or-later
 
-from spectre_core.post_processing.event_handler_register import event_handler_map
-from spectre_core.post_processing.base import BaseEventHandler
-from spectre_core.file_handlers.configs import CaptureConfig
+from ._register import event_handler_map
+from spectre_core.post_processing._base import BaseEventHandler
+from spectre_core.capture_configs import CaptureConfig, PNames
 from spectre_core.exceptions import EventHandlerNotFoundError
 
-def get_event_handler(event_handler_key: str) -> BaseEventHandler:
+
+def _get_event_handler(event_handler_key: str) -> BaseEventHandler:
     EventHandler = event_handler_map.get(event_handler_key)
     if EventHandler is None:
         valid_event_handler_keys = list(event_handler_map.keys())
@@ -18,5 +19,5 @@ def get_event_handler(event_handler_key: str) -> BaseEventHandler:
 
 def get_event_handler_from_tag(tag: str) -> BaseEventHandler:
     capture_config = CaptureConfig(tag)
-    event_handler_key = capture_config.get('event_handler_key')
-    return get_event_handler(event_handler_key)
+    event_handler_key = capture_config.get_parameter_value(PNames.EVENT_HANDLER_KEY)
+    return _get_event_handler(event_handler_key)

spectre_core/post_processing/{post_processor.py → _post_processor.py}

@@ -8,8 +8,8 @@ _LOGGER = getLogger(__name__)
 from watchdog.observers import Observer
 from watchdog.events import FileCreatedEvent
 
-from spectre_core.post_processing.factory import get_event_handler_from_tag
-from spectre_core.cfg import CHUNKS_DIR_PATH
+from ._factory import get_event_handler_from_tag
+from spectre_core.config import get_chunks_dir_path
 
 class PostProcessor:
     def __init__(self, 
@@ -24,7 +24,7 @@ class PostProcessor:
     def start(self):
         """Start an observer to process newly created files in the chunks directory"""
         self._observer.schedule(self._event_handler, 
-                                CHUNKS_DIR_PATH, 
+                                get_chunks_dir_path(), 
                                 recursive=True,
                                 event_filter=[FileCreatedEvent])