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

spectre_core/post_processing/_base.py

@@ -3,6 +3,7 @@
 # SPDX-License-Identifier: GPL-3.0-or-later
 
 from logging import getLogger
+
 _LOGGER = getLogger(__name__)
 
 from typing import Optional, cast
@@ -15,34 +16,26 @@ from spectre_core.capture_configs import CaptureConfig, PName
 from spectre_core.spectrograms import Spectrogram, join_spectrograms
 
 
-def make_sft_instance(
-    capture_config: CaptureConfig
-) -> ShortTimeFFT:
+def make_sft_instance(capture_config: CaptureConfig) -> ShortTimeFFT:
     """Extract window parameters from the input capture config and create an instance
     of `ShortTimeFFT` from `scipy.signal`.
 
     :param capture_config: The capture config storing the parameters.
-    :return: An instance of `ShortTimeFFT` consistent with the window parameters 
+    :return: An instance of `ShortTimeFFT` consistent with the window parameters
     in the capture config.
     """
-    window_type   = cast(str, capture_config.get_parameter_value(PName.WINDOW_TYPE))
-    sample_rate   = cast(int, capture_config.get_parameter_value(PName.SAMPLE_RATE))
-    window_hop    = cast(int, capture_config.get_parameter_value(PName.WINDOW_HOP))
-    window_size   = cast(int, capture_config.get_parameter_value(PName.WINDOW_SIZE))
-    window = get_window(window_type, 
-                        window_size)
-    return ShortTimeFFT(window, 
-                        window_hop,
-                        sample_rate, 
-                        fft_mode = "centered")
+    window_type = cast(str, capture_config.get_parameter_value(PName.WINDOW_TYPE))
+    sample_rate = cast(int, capture_config.get_parameter_value(PName.SAMPLE_RATE))
+    window_hop = cast(int, capture_config.get_parameter_value(PName.WINDOW_HOP))
+    window_size = cast(int, capture_config.get_parameter_value(PName.WINDOW_SIZE))
+    window = get_window(window_type, window_size)
+    return ShortTimeFFT(window, window_hop, sample_rate, fft_mode="centered")
 
 
 class BaseEventHandler(ABC, FileSystemEventHandler):
     """An abstract base class for event-driven file post-processing."""
-    def __init__(
-        self, 
-        tag: str
-    ) -> None:
+
+    def __init__(self, tag: str) -> None:
         """Initialise a `BaseEventHandler` instance.
 
         :param tag: The tag of the capture config used to capture the data.
@@ -51,7 +44,7 @@ class BaseEventHandler(ABC, FileSystemEventHandler):
 
         # load the capture config corresponding to the input tag
         self._capture_config = CaptureConfig(tag)
-        
+
         # store the next file to be processed (specifically, the absolute file path of the file)
         self._queued_file: Optional[str] = None
 
@@ -59,25 +52,17 @@ class BaseEventHandler(ABC, FileSystemEventHandler):
         # this can be flushed periodically to file as required.
         self._cached_spectrogram: Optional[Spectrogram] = None
 
-
     @abstractmethod
-    def process(
-        self, 
-        absolute_file_path: str
-    ) -> None:
+    def process(self, absolute_file_path: str) -> None:
         """
         Process a batch file at the given file path.
 
         :param absolute_file_path: The absolute path to the batch file to be processed.
         """
 
-
-    def on_created(
-        self, 
-        event: FileSystemEvent
-    ) -> None:
+    def on_created(self, event: FileSystemEvent) -> None:
         """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.
@@ -86,49 +71,49 @@ class BaseEventHandler(ABC, FileSystemEventHandler):
         """
         # The `src_path`` attribute holds the absolute path of the freshly closed file
         absolute_file_path = event.src_path
-        
+
         # Only process a file if:
         #
         # - It's extension matches the `watch_extension` as defined in the capture config.
         # - It's tag matches the current sessions tag.
         #
-        # This is important for two reasons. 
-        # 
-        # In the case of one session, the capture worker may write to two batch files simultaneously 
-        # (e.g., raw data file + seperate metadata file). We want to process them together - but this method will get called 
+        # This is important for two reasons.
+        #
+        # In the case of one session, the capture worker may write to two batch files simultaneously
+        # (e.g., raw data file + seperate metadata file). We want to process them together - but this method will get called
         # seperately for both file creation events. So, we filter by extension to account for this.
         #
         # Additionally in the case of multiple sessions, the capture workers will create batch files in the same directory concurrently.
         # This method is triggered for all file creation events, so we ensure the batch file tag matches the session tag and early return
         # otherwise. This way, each post processor worker picks up the right files to process.
-        watch_extension = cast(str, self._capture_config.get_parameter_value(PName.WATCH_EXTENSION))
+        watch_extension = cast(
+            str, self._capture_config.get_parameter_value(PName.WATCH_EXTENSION)
+        )
         if not absolute_file_path.endswith(f"{self._tag}.{watch_extension}"):
             return
-        
+
         _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)
+                _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:
+    def _cache_spectrogram(self, spectrogram: Spectrogram) -> None:
         """Cache the input spectrogram by storing it in the `_cached_spectrogram` attribute.
-        
+
         If the time range of the cached spectrogram exceeds that as specified in the capture config
         `PName.TIME_RANGE` parameter, the spectrogram in the cache is flushed to file. If `PName.TIME_RANGE`
         is nulled, the cache is flushed immediately.
@@ -140,20 +125,21 @@ class BaseEventHandler(ABC, FileSystemEventHandler):
         if self._cached_spectrogram is None:
             self._cached_spectrogram = spectrogram
         else:
-            self._cached_spectrogram = join_spectrograms([self._cached_spectrogram, spectrogram])
-  
+            self._cached_spectrogram = join_spectrograms(
+                [self._cached_spectrogram, spectrogram]
+            )
+
         time_range = self._capture_config.get_parameter_value(PName.TIME_RANGE) or 0.0
         if self._cached_spectrogram.time_range >= cast(float, time_range):
             self._flush_cache()
-    
 
-    def _flush_cache(
-        self
-    ) -> None:
+    def _flush_cache(self) -> None:
         """Flush the cached spectrogram to file."""
         if self._cached_spectrogram:
-            _LOGGER.info(f"Flushing spectrogram to file with start time "
-                         f"'{self._cached_spectrogram.format_start_time()}'")
+            _LOGGER.info(
+                f"Flushing spectrogram to file with start time "
+                f"'{self._cached_spectrogram.format_start_time()}'"
+            )
             self._cached_spectrogram.save()
             _LOGGER.info("Flush successful, resetting spectrogram cache")
-            self._cached_spectrogram = None # reset the cache
+            self._cached_spectrogram = None  # reset the cache

spectre_core/post_processing/_factory.py

@@ -10,8 +10,9 @@ from spectre_core.exceptions import EventHandlerNotFoundError
 from .plugins._event_handler_keys import EventHandlerKey
 from ._register import event_handler_map
 
+
 def _get_event_handler_cls_from_key(
-    event_handler_key: EventHandlerKey
+    event_handler_key: EventHandlerKey,
 ) -> Type[BaseEventHandler]:
     """Get a registered `BaseEventHandler` subclass.
 
@@ -22,14 +23,16 @@ def _get_event_handler_cls_from_key(
     event_handler_cls = event_handler_map.get(event_handler_key)
     if event_handler_cls is None:
         event_handler_keys = list(event_handler_map.keys())
-        raise EventHandlerNotFoundError((f"No event handler found for the capture mode '{event_handler_key}'. "
-                                         f"Please specify one of the following capture modes: {event_handler_keys}"))
+        raise EventHandlerNotFoundError(
+            (
+                f"No event handler found for the capture mode '{event_handler_key}'. "
+                f"Please specify one of the following capture modes: {event_handler_keys}"
+            )
+        )
     return event_handler_cls
 
 
-def get_event_handler_cls_from_tag(
-    tag: str
-) -> Type[BaseEventHandler]:
+def get_event_handler_cls_from_tag(tag: str) -> Type[BaseEventHandler]:
     """
     Retrieve the event handler class, using the event handler key stored in a capture config.
 
@@ -37,13 +40,13 @@ def get_event_handler_cls_from_tag(
     :return: The event handler class specified in the capture config.
     """
     capture_config = CaptureConfig(tag)
-    event_handler_key = cast(str, capture_config.get_parameter_value(PName.EVENT_HANDLER_KEY))
-    return _get_event_handler_cls_from_key( EventHandlerKey(event_handler_key) )
+    event_handler_key = cast(
+        str, capture_config.get_parameter_value(PName.EVENT_HANDLER_KEY)
+    )
+    return _get_event_handler_cls_from_key(EventHandlerKey(event_handler_key))
 
 
-def get_event_handler(
-    tag: str
-) -> BaseEventHandler:
+def get_event_handler(tag: str) -> BaseEventHandler:
     """Create an event handler class instance, using the event handler key stored in a capture config.
 
     :param tag: The capture config tag.

spectre_core/post_processing/_post_processor.py

@@ -3,6 +3,7 @@
 # SPDX-License-Identifier: GPL-3.0-or-later
 
 from logging import getLogger
+
 _LOGGER = getLogger(__name__)
 
 from watchdog.observers import Observer
@@ -12,27 +13,30 @@ from ._factory import get_event_handler
 from spectre_core.config import get_batches_dir_path
 
 
-def start_post_processor(
-    tag: str
-) -> None:
+def start_post_processor(tag: str) -> None:
     """Start a thread to process newly created files in the `batches` directory.
 
     :param tag: The tag of the capture config used for data capture.
     """
     post_processor = Observer()
     event_handler = get_event_handler(tag)
-    post_processor.schedule(event_handler, 
-                            get_batches_dir_path(), 
-                            recursive=True,
-                            event_filter=[ FileCreatedEvent ])
-    
+    post_processor.schedule(
+        event_handler,
+        get_batches_dir_path(),
+        recursive=True,
+        event_filter=[FileCreatedEvent],
+    )
+
     try:
-        _LOGGER.info("Starting the post processing thread...") 
+        _LOGGER.info("Starting the post processing thread...")
         post_processor.start()
         post_processor.join()
     except KeyboardInterrupt:
-        _LOGGER.warning(("Keyboard interrupt detected. Signalling "
-                         "the post processing thread to stop"))
+        _LOGGER.warning(
+            (
+                "Keyboard interrupt detected. Signalling "
+                "the post processing thread to stop"
+            )
+        )
         post_processor.stop()
         _LOGGER.warning(("Post processing thread has been successfully stopped"))
-

spectre_core/post_processing/_register.py

@@ -10,9 +10,11 @@ from ._base import BaseEventHandler
 # Map populated at runtime via the `register_event_handler` decorator.
 event_handler_map: dict[EventHandlerKey, Type[BaseEventHandler]] = {}
 
-T = TypeVar('T', bound=BaseEventHandler)
+T = TypeVar("T", bound=BaseEventHandler)
+
+
 def register_event_handler(
-    event_handler_key: EventHandlerKey
+    event_handler_key: EventHandlerKey,
 ) -> Callable[[Type[T]], Type[T]]:
     """Decorator to register a `BaseEventHandler` subclass under a specified `EventHandlerKey`.
 
@@ -20,12 +22,13 @@ def register_event_handler(
     :raises ValueError: If the provided `event_handler_key` is already registered.
     :return: A decorator that registers the `BaseBatch` subclass under the given `batch_key`.
     """
-    def decorator(
-        cls: Type[T]
-    ) -> Type[T]:
+
+    def decorator(cls: Type[T]) -> Type[T]:
         if event_handler_key in event_handler_map:
-            raise ValueError(f"An event handler with key '{event_handler_key}' is already registered!")
+            raise ValueError(
+                f"An event handler with key '{event_handler_key}' is already registered!"
+            )
         event_handler_map[event_handler_key] = cls
         return cls
-    return decorator
 
+    return decorator

spectre_core/post_processing/plugins/_event_handler_keys.py

@@ -4,13 +4,14 @@
 
 from enum import Enum
 
-class EventHandlerKey(Enum):  
+
+class EventHandlerKey(Enum):
     """Key bound to a `BaseEventHandler` plugin class.
-    
+
     :ivar FIXED_CENTER_FREQUENCY: Postprocess data capture at a fixed center frequency.
     :ivar SWEPT_CENTER_FREQUENCY: Postprocess data capture where the center frequency is continually sweeping
     in fixed increments.
     """
+
     FIXED_CENTER_FREQUENCY = "fixed_center_frequency"
     SWEPT_CENTER_FREQUENCY = "swept_center_frequency"
-    

spectre_core/post_processing/plugins/_fixed_center_frequency.py

@@ -3,6 +3,7 @@
 # SPDX-License-Identifier: GPL-3.0-or-later
 
 from logging import getLogger
+
 _LOGGER = getLogger(__name__)
 
 import numpy as np
@@ -15,7 +16,12 @@ import numpy as np
 
 from spectre_core.capture_configs import CaptureConfig, PName
 from spectre_core.batches import IQStreamBatch
-from spectre_core.spectrograms import Spectrogram, SpectrumUnit, time_average, frequency_average
+from spectre_core.spectrograms import (
+    Spectrogram,
+    SpectrumUnit,
+    time_average,
+    frequency_average,
+)
 from ._event_handler_keys import EventHandlerKey
 from .._base import BaseEventHandler, make_sft_instance
 from .._register import register_event_handler
@@ -26,11 +32,11 @@ def _do_stfft(
     capture_config: CaptureConfig,
 ) -> Tuple[npt.NDArray[np.float32], npt.NDArray[np.float32], npt.NDArray[np.float32]]:
     """Do a Short-time Fast Fourier Transform on an array of complex IQ samples.
-    
-    The computation requires extra metadata, which is extracted from the detached header in the batch 
+
+    The computation requires extra metadata, which is extracted from the detached header in the batch
     and the capture config used to capture the data.
-    
-    The current implementation relies heavily on the `ShortTimeFFT` implementation from 
+
+    The current implementation relies heavily on the `ShortTimeFFT` implementation from
     `scipy.signal` (https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.ShortTimeFFT.html)
     which takes up a lot of the compute time.
     """
@@ -38,64 +44,64 @@ def _do_stfft(
     sft = make_sft_instance(capture_config)
 
     # set p0=0, since by convention in the STFFT docs, p=0 corresponds to the slice centred at t=0
-    p0=0
+    p0 = 0
 
     # set p1 to the index of the first slice where the "midpoint" of the window is still inside the signal
     num_samples = len(iq_data)
     p1 = sft.upper_border_begin(num_samples)[1]
-    
+
     # compute a ShortTimeFFT on the IQ samples
-    complex_spectra = sft.stft(iq_data, 
-                               p0 = p0, 
-                               p1 = p1) 
-    
+    complex_spectra = sft.stft(iq_data, p0=p0, p1=p1)
+
     # compute the magnitude of each spectral component
     dynamic_spectra = np.abs(complex_spectra)
 
-
     # assign a physical time to each spectrum
     # p0 is defined to correspond with the first sample, at t=0 [s]
-    times = sft.t(num_samples, 
-                  p0 = p0, 
-                  p1 = p1)
+    times = sft.t(num_samples, p0=p0, p1=p1)
     # assign physical frequencies to each spectral component
-    frequencies = sft.f + cast(float, capture_config.get_parameter_value(PName.CENTER_FREQUENCY))
+    frequencies = sft.f + cast(
+        float, capture_config.get_parameter_value(PName.CENTER_FREQUENCY)
+    )
 
-    return times.astype(np.float32), frequencies.astype(np.float32), dynamic_spectra.astype(np.float32)
+    return (
+        times.astype(np.float32),
+        frequencies.astype(np.float32),
+        dynamic_spectra.astype(np.float32),
+    )
 
 
 def _build_spectrogram(
-    batch: IQStreamBatch,
-    capture_config: CaptureConfig
+    batch: IQStreamBatch, capture_config: CaptureConfig
 ) -> Spectrogram:
     """Generate a spectrogram using `IQStreamBatch` IQ samples."""
     # read the data from the batch
     iq_metadata = batch.hdr_file.read()
-    iq_samples  = batch.bin_file.read()
+    iq_samples = batch.bin_file.read()
 
-    times, frequencies, dynamic_spectra = _do_stfft(iq_samples,
-                                                    capture_config)
+    times, frequencies, dynamic_spectra = _do_stfft(iq_samples, capture_config)
 
     # compute the start datetime for the spectrogram by adding the millisecond component to the batch start time
-    spectrogram_start_datetime = batch.start_datetime + timedelta(milliseconds=iq_metadata.millisecond_correction)
-    
-    return Spectrogram(dynamic_spectra,
-                       times,
-                       frequencies,
-                       batch.tag,
-                       SpectrumUnit.AMPLITUDE,
-                       spectrogram_start_datetime)
+    spectrogram_start_datetime = batch.start_datetime + timedelta(
+        milliseconds=iq_metadata.millisecond_correction
+    )
+
+    return Spectrogram(
+        dynamic_spectra,
+        times,
+        frequencies,
+        batch.tag,
+        SpectrumUnit.AMPLITUDE,
+        spectrogram_start_datetime,
+    )
 
 
 @register_event_handler(EventHandlerKey.FIXED_CENTER_FREQUENCY)
 class FixedEventHandler(BaseEventHandler):
-    def process(
-        self, 
-        absolute_file_path: str
-    ) -> None:
+    def process(self, absolute_file_path: str) -> None:
         """Compute a spectrogram using `IQStreamBatch` IQ samples, cache it, then save it to file in the FITS
         format. The IQ samples are assumed to have been collected at a fixed center frequency.
-        
+
         The computed spectrogram is averaged in time and frequency as per the user-configured capture config.
         Once the spectrogram has been computed successfully, the `.bin` and `.hdr` files are removed.
 
@@ -104,22 +110,23 @@ class FixedEventHandler(BaseEventHandler):
         _LOGGER.info(f"Processing: {absolute_file_path}")
         file_name = os.path.basename(absolute_file_path)
         base_file_name, _ = os.path.splitext(file_name)
-        batch_start_time, tag = base_file_name.split('_')
+        batch_start_time, tag = base_file_name.split("_")
 
         batch = IQStreamBatch(batch_start_time, tag)
 
         _LOGGER.info("Creating spectrogram")
-        spectrogram = _build_spectrogram(batch,
-                                         self._capture_config)
-        
-        time_resolution = cast(float, self._capture_config.get_parameter_value(PName.TIME_RESOLUTION))
-        spectrogram = time_average(spectrogram,
-                                   resolution=time_resolution)
-
-        frequency_resolution = cast(float, self._capture_config.get_parameter_value(PName.FREQUENCY_RESOLUTION))
-        spectrogram = frequency_average(spectrogram,
-                                        resolution=frequency_resolution)
-        
+        spectrogram = _build_spectrogram(batch, self._capture_config)
+
+        time_resolution = cast(
+            float, self._capture_config.get_parameter_value(PName.TIME_RESOLUTION)
+        )
+        spectrogram = time_average(spectrogram, resolution=time_resolution)
+
+        frequency_resolution = cast(
+            float, self._capture_config.get_parameter_value(PName.FREQUENCY_RESOLUTION)
+        )
+        spectrogram = frequency_average(spectrogram, resolution=frequency_resolution)
+
         self._cache_spectrogram(spectrogram)
 
         _LOGGER.info(f"Deleting {batch.bin_file.file_path}")