spectre-core 0.0.12__py3-none-any.whl → 0.0.13__py3-none-any.whl

spectre_core/post_processing/_base.py

@@ -5,23 +5,30 @@
 from logging import getLogger
 _LOGGER = getLogger(__name__)
 
-from typing import Optional
+from typing import Optional, cast
 from abc import ABC, abstractmethod
 from scipy.signal import ShortTimeFFT, get_window
 
-from watchdog.events import FileSystemEventHandler, FileCreatedEvent
+from watchdog.events import FileSystemEventHandler, FileSystemEvent
 
-from spectre_core.capture_configs import CaptureConfig, PNames
-from spectre_core.batches import BaseBatch, get_batch_cls_from_tag
+from spectre_core.capture_configs import CaptureConfig, PName
 from spectre_core.spectrograms import Spectrogram, join_spectrograms
 
 
-def make_sft_instance(capture_config: CaptureConfig
+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)
+    """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 
+    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, 
@@ -31,52 +38,60 @@ def make_sft_instance(capture_config: CaptureConfig
 
 
 class BaseEventHandler(ABC, FileSystemEventHandler):
-    def __init__(self, 
-                 tag: str):
+    """An abstract base class for event-driven file post-processing."""
+    def __init__(
+        self, 
+        tag: str
+    ) -> None:
+        """Initialise a `BaseEventHandler` instance.
+
+        :param tag: The tag of the capture config used to capture the data.
+        """
         self._tag = tag
 
-        # the tag tells us 'what type' of data is stored in the files for each batch
-        self._Batch = get_batch_cls_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)
-
+        # 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
 
-        # store batched spectrograms as they are created into a cache
-        # which is flushed periodically according to a user defined 
-        # time range
+        # optionally store batched spectrograms as they are created into a cache
+        # this can be flushed periodically to file as required.
         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 process(
+        self, 
+        absolute_file_path: str
+    ) -> None:
         """
+        Process a batch file at the given file path.
 
-    def on_created(self, 
-                   event: FileCreatedEvent):
+        :param absolute_file_path: The absolute path to the batch file to be processed.
+        """
+
+
+    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.
-        """
 
-        # the 'src_path' attribute holds the absolute path of the newly created file
+        :param event: The file system event containing the file details.
+        """
+        # the `src_path`` attribute holds the absolute path of the freshly closed 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):
+        # only 'notice' a file if it ends with the appropriate extension as defined in the capture config
+        watch_extension = cast(str, self._capture_config.get_parameter_value(PName.WATCH_EXTENSION))
+        
+        if absolute_file_path.endswith( 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:
@@ -92,28 +107,39 @@ class BaseEventHandler(ABC, FileSystemEventHandler):
             # 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.
+
+        :param spectrogram: The spectrogram to store in the cache.
+        """
         _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 0.0
   
-        if self._cached_spectrogram.time_range >= time_range:
+        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(precise=True)}'")
+                         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

spectre_core/post_processing/_factory.py

@@ -2,22 +2,52 @@
 # This file is part of SPECTRE
 # SPDX-License-Identifier: GPL-3.0-or-later
 
-from ._register import event_handler_map
+from typing import Type, cast
+
 from spectre_core.post_processing._base import BaseEventHandler
-from spectre_core.capture_configs import CaptureConfig, PNames
+from spectre_core.capture_configs import CaptureConfig, PName
 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
+) -> Type[BaseEventHandler]:
+    """Get a registered `BaseEventHandler` subclass.
 
+    :param event_handler_key: The key used to register the `BaseEventHandler` subclass.
+    :raises EventHandlerNotFoundError: If an invalid `event_handler_key` is provided.
+    :return: The `BaseEventHandler` subclass registered under `event_handler_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}"))
+    return event_handler_cls
 
-def _get_event_handler_cls(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())
-        raise EventHandlerNotFoundError((f"No event handler found for the event handler key '{event_handler_key}'. "
-                                         f"Please specify one of the following event handler keys: {valid_event_handler_keys}"))
-    return EventHandler
 
+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.
 
-def get_event_handler_cls_from_tag(tag: str) -> BaseEventHandler:
+    :param tag: The capture config tag.
+    :return: The event handler class specified in the capture config.
+    """
     capture_config = CaptureConfig(tag)
-    event_handler_key = capture_config.get_parameter_value(PNames.EVENT_HANDLER_KEY)
-    return _get_event_handler_cls(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:
+    """Create an event handler class instance, using the event handler key stored in a capture config.
+
+    :param tag: The capture config tag.
+    :return: An instance of the event handler class.
+    """
+    event_handler_cls = get_event_handler_cls_from_tag(tag)
+    return event_handler_cls(tag)

spectre_core/post_processing/_post_processor.py

@@ -8,33 +8,31 @@ _LOGGER = getLogger(__name__)
 from watchdog.observers import Observer
 from watchdog.events import FileCreatedEvent
 
-from ._factory import get_event_handler_cls_from_tag
+from ._factory import get_event_handler
 from spectre_core.config import get_batches_dir_path
 
-class PostProcessor:
-    def __init__(self, 
-                 tag: str):
-        
-        self._observer = Observer()
 
-        EventHandler = get_event_handler_cls_from_tag(tag)
-        self._event_handler = EventHandler(tag)
-
-
-    def start(self):
-        """Start an observer to process newly created files in the batches directory"""
-        self._observer.schedule(self._event_handler, 
-                                get_batches_dir_path(), 
-                                recursive=True,
-                                event_filter=[FileCreatedEvent])
-        
-        try:
-            _LOGGER.info("Starting the post processing thread...") 
-            self._observer.start()
-            self._observer.join()
-        except KeyboardInterrupt:
-            _LOGGER.warning(("Keyboard interrupt detected. Signalling "
-                             "the post processing thread to stop"))
-            self._observer.stop()
-            _LOGGER.warning(("Post processing thread has been successfully stopped"))
+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 ])
+    
+    try:
+        _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"))
+        post_processor.stop()
+        _LOGGER.warning(("Post processing thread has been successfully stopped"))
 

spectre_core/post_processing/_register.py

@@ -2,13 +2,29 @@
 # This file is part of SPECTRE
 # SPDX-License-Identifier: GPL-3.0-or-later
 
-# Global dictionaries to hold the mappings
-event_handler_map = {}
+from typing import Type, Callable, TypeVar
 
-# classes decorated with @register_event_handler([EVENT_HANDLER_KEY])
-# will be added to event_handler_map
-def register_event_handler(event_handler_key: str):
-    def decorator(cls):
+from .plugins._event_handler_keys import EventHandlerKey
+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)
+def register_event_handler(
+    event_handler_key: EventHandlerKey
+) -> Callable[[Type[T]], Type[T]]:
+    """Decorator to register a `BaseEventHandler` subclass under a specified `EventHandlerKey`.
+
+    :param event_handler_key: The key to register the `BaseEventHandler` subclass.
+    :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]:
+        if event_handler_key in event_handler_map:
+            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

spectre_core/post_processing/plugins/_event_handler_keys.py

@@ -0,0 +1,16 @@
+# SPDX-FileCopyrightText: © 2024 Jimmy Fitzpatrick <jcfitzpatrick12@gmail.com>
+# This file is part of SPECTRE
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+from enum import 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

@@ -0,0 +1,129 @@
+# 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__)
+
+import numpy as np
+from typing import Tuple, cast
+from datetime import timedelta
+
+import os
+import numpy.typing as npt
+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 ._event_handler_keys import EventHandlerKey
+from .._base import BaseEventHandler, make_sft_instance
+from .._register import register_event_handler
+
+
+def _do_stfft(
+    iq_data: npt.NDArray[np.complex64],
+    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 
+    and the capture config used to capture the data.
+    
+    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.
+    """
+
+    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
+
+    # 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) 
+    
+    # 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)
+    # assign physical frequencies to each spectral component
+    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)
+
+
+def _build_spectrogram(
+    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()
+
+    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)
+
+
+@register_event_handler(EventHandlerKey.FIXED_CENTER_FREQUENCY)
+class FixedEventHandler(BaseEventHandler):
+    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.
+
+        :param absolute_file_path: The absolute file path of the `.bin` file in the batch.
+        """
+        _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 = 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)
+        
+        self._cache_spectrogram(spectrogram)
+
+        _LOGGER.info(f"Deleting {batch.bin_file.file_path}")
+        batch.bin_file.delete()
+
+        _LOGGER.info(f"Deleting {batch.hdr_file.file_path}")
+        batch.hdr_file.delete()