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

spectre_core/batches/_batches.py

@@ -2,97 +2,142 @@
 # This file is part of SPECTRE
 # SPDX-License-Identifier: GPL-3.0-or-later
 
-from logging import getLogger
-_LOGGER = getLogger(__name__)
-
 import os
-from typing import Optional
+from typing import Optional, TypeVar, Type, Generic, Iterator
 from collections import OrderedDict
-import warnings
 from datetime import datetime
 
+from spectre_core.config import TimeFormat
 from spectre_core.spectrograms import Spectrogram, time_chop, join_spectrograms
-from spectre_core.config import get_batches_dir_path, TimeFormats
+from spectre_core.config import get_batches_dir_path
 from spectre_core.exceptions import (
-    SpectrogramNotFoundError,
     BatchNotFoundError
 )
 from ._base import BaseBatch
-from ._factory import get_batch_cls_from_tag
-
-class Batches:
-    """A collection of batches for a given day of the year."""
-    def __init__(self, 
-                 tag: str,
-                 year: Optional[int] = None, 
-                 month: Optional[int] = None, 
-                 day: Optional[int] = None):
+
+T = TypeVar('T', bound=BaseBatch)
+class Batches(Generic[T]):
+    """Managed collection of `Batch` instances for a given tag. Provides a simple
+    interface for read operations on batched data files."""
+    def __init__(
+        self, 
+        tag: str,
+        batch_cls: Type[T],
+        year: Optional[int] = None, 
+        month: Optional[int] = None, 
+        day: Optional[int] = None
+    ) -> None:
+        """Initialise a `Batches` instance.
+
+        :param tag: The batch name tag.
+        :param batch_cls: The `Batch` class used to read data files tagged by `tag`.
+        :param year: Filter batch files under a numeric year. Defaults to None.
+        :param month: Filter batch files under a numeric month. Defaults to None.
+        :param day: Filter batch files under a numeric day. Defaults to None.
+        """
         self._tag = tag
-        self._Batch = get_batch_cls_from_tag(tag)
-        self._batch_map: dict[str, BaseBatch] = OrderedDict()
+        self._batch_cls = batch_cls
+        self._batch_map: dict[str, T] = OrderedDict()
         self.set_date(year, month, day)
 
 
     @property
-    def tag(self) -> str:
-        """Tag identifier for each batch."""
+    def tag(
+        self
+    ) -> str:
+        """The batch name tag."""
         return self._tag
 
 
     @property
-    def year(self) -> int:
-        """The numeric year."""
+    def batch_cls(
+        self
+    ) -> Type[T]:
+        """The `Batch` class used to read the batched files."""
+        return self._batch_cls
+
+
+    @property
+    def year(
+        self
+    ) -> Optional[int]:
+        """The numeric year, to filter batch files."""
         return self._year
 
 
     @property 
-    def month(self) -> int:
-        """The numeric month of the year."""
+    def month(
+        self
+    ) -> Optional[int]:
+        """The numeric month of the year, to filter batch files."""
         return self._month
     
 
     @property
-    def day(self) -> int:
-        """The numeric day of the year."""
+    def day(
+        self
+    ) -> Optional[int]:
+        """The numeric day of the year, to filter batch files."""
         return self._day
     
 
     @property
-    def batches_dir_path(self) -> str:
-        """The parent directory for all the batches."""
+    def batches_dir_path(
+        self
+    ) -> str:
+        """The shared ancestral path for all the batches. `Batches` recursively searches
+        this directory to find all batches whose batch name contains `tag`."""
         return get_batches_dir_path(self.year, self.month, self.day)
     
 
     @property
-    def batch_list(self) -> list[BaseBatch]:
-        """A list of all the batch instances."""
+    def batch_list(
+        self
+    ) -> list[T]:
+        """A list of all batches found within `batches_dir_path`."""
         return  list(self._batch_map.values())
     
 
     @property
-    def start_times(self) -> list[str]:
-        """The start times of each batch."""
+    def start_times(
+        self
+    ) -> list[str]:
+        """The start times of each batch found within `batches_dir_path`."""
         return list(self._batch_map.keys())
 
 
     @property
-    def num_batches(self) -> int:
-        """The number of batches in the batch parent directory."""
+    def num_batches(
+        self
+    ) -> int:
+        """The total number of batches found within `batches_dir_path`."""
         return len(self.batch_list)
 
 
-    def set_date(self, 
-                 year: Optional[int],
-                 month: Optional[int],
-                 day: Optional[int]) -> None:
-        """Update the parent directory for the batches according to the numeric date."""
+    def set_date(
+        self, 
+        year: Optional[int],
+        month: Optional[int],
+        day: Optional[int]
+    ) -> None:
+        """Reset `batches_dir_path` according to the numeric date, and refresh the list
+        of available batches.
+
+        :param year: Filter by the numeric year.
+        :param month: Filter by the numeric month of the year.
+        :param day: Filter by the numeric day of the month.
+        """
         self._year = year
         self._month = month
         self._day = day
-        self._update_batch_map()
+        self.update()
 
 
-    def _update_batch_map(self) -> None:
+    def update(
+        self
+    ) -> None:
+        """Perform a fresh search all files in `batches_dir_path` for batches 
+        with `tag` in the batch name."""
         # reset cache
         self._batch_map = OrderedDict() 
         
@@ -103,95 +148,100 @@ class Batches:
             batch_name, _ = os.path.splitext(batch_file_name)
             start_time, tag = batch_name.split("_", 1)
             if tag == self._tag:
-                self._batch_map[start_time] = self._Batch(start_time, tag)
+                self._batch_map[start_time] = self.batch_cls(start_time, tag)
         
         self._batch_map = OrderedDict(sorted(self._batch_map.items()))
-
-
-    def update(self) -> None:
-        """Public alias for setting batch map"""
-        self._update_batch_map()
     
 
-    def __iter__(self):
+    def __iter__(
+        self
+    ) -> Iterator[T]:
         """Iterate over the stored batch instances."""
         yield from self.batch_list
+        
+    
+    def __len__(
+        self
+    ):
+        return self.num_batches
 
 
-    def _get_from_start_time(self, 
-                             start_time: str) -> BaseBatch:
-        """Get the batch according to the input start time."""
+    def _get_from_start_time(
+        self, 
+        start_time: str
+    ) -> T:
+        """Find and return the `Batch` instance based on the string formatted start time."""
         try:
             return self._batch_map[start_time]
         except KeyError:
             raise BatchNotFoundError(f"Batch with start time {start_time} could not be found within {self.batches_dir_path}")
 
 
-    def _get_from_index(self, 
-                        index: int) -> BaseBatch:
-        """Get the batch according to its index, where the batches are ordered in time."""
-        num_batches = len(self.batch_list)
-        if num_batches == 0:
+    def _get_from_index(
+        self, 
+        index: int
+    ) -> T:
+        """Find and return the `Batch` instance based on its numeric index.
+        
+        Batches are ordered sequentially in time, so index `0` corresponds to the oldest
+        `Batch` with respect to the start time.
+        """
+        if self.num_batches == 0:
             raise BatchNotFoundError("No batches are available")
-        index = index % num_batches  # Use modulo to make the index wrap around. Allows the user to iterate over all the batches via index cyclically.
+        elif index > self.num_batches:
+            raise IndexError(f"Index '{index}' is greater than the number of batches '{self.num_batches}'")
         return self.batch_list[index]
 
 
-    def __getitem__(self, subscript: str | int):
+    def __getitem__(
+        self, 
+        subscript: str | int
+    ) -> T:
+        """Get a `Batch` instanced based on either the start time or chronological index.
+
+        :param subscript: If the subscript is a string, interpreted as a formatted start time. 
+        If the subscript is an integer, it is interpreted as a chronological index.
+        :return: The corresponding `BaseBatch` subclass.
+        """
         if isinstance(subscript, str):
             return self._get_from_start_time(subscript)
         elif isinstance(subscript, int):
             return self._get_from_index(subscript)
-    
 
-    def num_batch_files(self, 
-                        extension: str) -> int:
-        """Get the number of existing batch files with the given extension."""
-        return sum(1 for batch_file in self if batch_file.has_file(extension))
 
+    def get_spectrogram_from_range(
+        self,
+        start_time: str, 
+        end_time: str
+    ) -> Spectrogram:
+        """
+        Retrieve a spectrogram spanning the specified time range. 
 
-    def get_spectrogram_from_range(self, 
-                                   start_time: str, 
-                                   end_time: str) -> Spectrogram:
-        """Return a spectrogram over the input time range."""
+        :param start_time: The start time of the range (inclusive).
+        :param end_time: The end time of the range (inclusive).
+        :raises FileNotFoundError: If no spectrogram data is available within the specified time range.
+        :return: A spectrogram created by stitching together data from all matching batches.
+        """
         # Convert input strings to datetime objects
-        start_datetime = datetime.strptime(start_time, TimeFormats.DATETIME)
-        end_datetime   = datetime.strptime(end_time, TimeFormats.DATETIME)
-
-        if start_datetime.day != end_datetime.day:
-            warning_message = "Joining spectrograms across multiple days"
-            _LOGGER.warning(warning_message)
-            warnings.warn(warning_message, RuntimeWarning)
+        start_datetime = datetime.strptime(start_time, TimeFormat.DATETIME)
+        end_datetime   = datetime.strptime(end_time,   TimeFormat.DATETIME)
 
         spectrograms = []
-        num_fits_batch_files = self.num_batch_files("fits")
-
-        for i, batch in enumerate(self):
-            # skip batches without fits files
-            if not batch.has_file("fits"):
+        for batch in self:
+            # skip batches without spectrogram data
+            if not batch.spectrogram_file.exists:
                 continue
-            
-            # rather than reading all files to evaluate the actual upper bound to their time range (slow)
-            # place an upper bound by using the start datetime for the next batch
-            # this assumes that the batches are non-overlapping (reasonable assumption)
-            lower_bound = batch.start_datetime
-            if i < num_fits_batch_files:
-                next_batch = self[i + 1]
-                upper_bound = next_batch.start_datetime
-            # if there is no "next batch" then we do have to read the file
-            else:
-                fits_batch = batch.get_file("fits")
-                upper_bound = fits_batch.datetimes[-1]
-
-            # if the batch overlaps with the input time range, then read the fits file
+
+            spectrogram = batch.read_spectrogram()
+            lower_bound = spectrogram.datetimes[0]
+            upper_bound = spectrogram.datetimes[-1]
+
+            # Check if the batch overlaps with the input time range
             if start_datetime <= upper_bound and lower_bound <= end_datetime:
-                spectrogram = batch.read_file("fits")
-                spectrogram = time_chop(spectrogram, start_time, end_time)
-                # if we have a non-empty spectrogram, append it to the list of spectrograms
-                if spectrogram:
-                    spectrograms.append(spectrogram)
+                spectrograms.append( time_chop(spectrogram, start_time, end_time) )
 
         if spectrograms:
             return join_spectrograms(spectrograms)
         else:
-            raise SpectrogramNotFoundError("No spectrogram data found for the given time range")
+            raise FileNotFoundError(f"No spectrogram data found for the time range "
+                                    f"{start_time} to {end_time}.")

spectre_core/batches/_factory.py

@@ -2,26 +2,68 @@
 # This file is part of SPECTRE
 # SPDX-License-Identifier: GPL-3.0-or-later
 
-from spectre_core.capture_configs import CaptureConfig, PNames
+from typing import Literal, overload, Type, cast
+
 from spectre_core.exceptions import BatchNotFoundError
-from ._register import batch_map
+from spectre_core.capture_configs import CaptureConfig, PName
 from ._base import BaseBatch
+from ._register import batch_map
+from .plugins._batch_keys import BatchKey
+from .plugins._callisto import CallistoBatch
+from .plugins._iq_stream import IQStreamBatch
+    
+
+@overload
+def get_batch_cls(
+    batch_key: Literal[BatchKey.CALLISTO],
+) -> Type[CallistoBatch]:
+    ...
 
 
-def _get_batch_cls(batch_key: str) -> BaseBatch:
-    Batch = batch_map.get(batch_key)
-    if Batch is None:
+@overload
+def get_batch_cls(
+    batch_key: Literal[BatchKey.IQ_STREAM],
+) -> Type[IQStreamBatch]:
+    ...
+    
+    
+@overload
+def get_batch_cls(batch_key: BatchKey) -> Type[BaseBatch]:
+    ...
+
+
+def get_batch_cls(
+    batch_key: BatchKey,
+) -> Type[BaseBatch]:
+    """Get a registered `BaseBatch` subclass.
+
+    :param batch_key: The key used to register the `BaseBatch` subclass.
+    :raises BatchNotFoundError: If an invalid `batch_key` is provided.
+    :return: The `BaseBatch` subclass corresponding to the input key.
+    """
+    batch_cls = batch_map.get(batch_key)
+    if batch_cls is None:
         valid_batch_keys = list(batch_map.keys())
-        raise BatchNotFoundError(f"No batch found for the batch key: {batch_key}. Valid batch keys are: {valid_batch_keys}")
-    return Batch
+        raise BatchNotFoundError(f"No batch found for the batch key: {batch_key}. "
+                                 f"Valid batch keys are: {valid_batch_keys}")
+    return batch_cls
 
 
-def get_batch_cls_from_tag(tag: str) -> BaseBatch:
-    # if we are dealing with a callisto batch, the batch key is equal to the tag
+def get_batch_cls_from_tag(
+    tag: str
+) -> Type[BaseBatch]:
+    # if the tag is reserved (i.e., corresponds to third-party spectrogram data)
+    # directly fetch the right class.
     if "callisto" in tag:
-        batch_key = "callisto"
-    # otherwise, we fetch the batch key from the capture config
+        return get_batch_cls(BatchKey.CALLISTO)
+
+    # otherwise, assume that the tag has an associated capture config,
     else:
-        capture_config= CaptureConfig(tag)
-        batch_key = capture_config.get_parameter_value(PNames.BATCH_KEY)
-    return _get_batch_cls(batch_key)
+        capture_config = CaptureConfig(tag)
+        if PName.BATCH_KEY not in capture_config.parameters.name_list:
+            raise ValueError(f"Could not infer batch class from the tag 'tag'. "
+                             f"A parameter with name `{PName.BATCH_KEY.value}` "
+                             f"does not exist.")
+        
+        batch_key = BatchKey( cast(str, capture_config.get_parameter_value( PName.BATCH_KEY) ) )
+        return get_batch_cls( batch_key )

spectre_core/batches/_register.py

@@ -2,14 +2,29 @@
 # This file is part of SPECTRE
 # SPDX-License-Identifier: GPL-3.0-or-later
 
-#  Global dictionaries to hold the mappings
-batch_map = {}
+from typing import Type, Callable, TypeVar
 
-# classes decorated with @register_batch([BATCH_KEY])
-# will be added to batch_map
-def register_batch(batch_key: str):
-    def decorator(cls):
+from ._base import BaseBatch
+from .plugins._batch_keys import BatchKey
+
+# Map populated at runtime via the `register_batch` decorator.
+batch_map: dict[BatchKey, Type[BaseBatch]] = {}
+
+T = TypeVar('T', bound=BaseBatch)
+def register_batch(
+    batch_key: BatchKey
+) -> Callable[[Type[T]], Type[T]]:
+    """Decorator to register a `BaseBatch` subclass under a specified `BatchKey`.
+
+    :param batch_key: The key to register the `BaseBatch` subclass under.
+    :raises ValueError: If the provided `batch_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 batch_key in batch_map:
+            raise ValueError(f"A batch with key '{batch_key}' is already registered!")
         batch_map[batch_key] = cls
         return cls
-    return decorator
-
+    return decorator

spectre_core/batches/plugins/_batch_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 BatchKey(Enum):  
+    """Key bound to a `Batch` plugin class.
+    
+    :ivar IQ_STREAM: Represents the default batch data generated  by `spectre`, 
+    containing IQ stream data and other data derived from it.
+    :ivar CALLISTO: Represents FITS files generated by the e-Callisto network.
+    """
+    IQ_STREAM  = "iq_stream"
+    CALLISTO   = "callisto"
+    

spectre_core/batches/plugins/_callisto.py

@@ -0,0 +1,183 @@
+# SPDX-FileCopyrightText: © 2024 Jimmy Fitzpatrick <jcfitzpatrick12@gmail.com>
+# This file is part of SPECTRE
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+from dataclasses import dataclass
+from datetime import datetime
+
+import numpy as np
+import numpy.typing as npt
+from astropy.io import fits
+from astropy.io.fits.hdu.image import PrimaryHDU
+from astropy.io.fits.hdu.table import BinTableHDU
+from astropy.io.fits.hdu.hdulist import HDUList
+
+from spectre_core.spectrograms import Spectrogram, SpectrumUnit
+from ._batch_keys import BatchKey
+from .._base import BaseBatch, BatchFile
+from .._register import register_batch
+
+
+@dataclass(frozen=True)
+class _BatchExtension:
+    """Supported extensions for a `CallistoBatch`.
+    
+    :ivar FITS: Corresponds to the `.fits` file extension.
+    """
+    FITS: str = "fits"
+    
+        
+class _FitsFile(BatchFile[Spectrogram]):
+    """Stores spectrogram data in the FITS format generated by the e-Callisto network."""
+    def __init__(
+        self, 
+        parent_dir_path: str, 
+        base_file_name: str
+    ) -> None:
+        """Initialise a `_FitsFile` instance.
+
+        :param parent_dir_path: The parent directory for the batch.
+        :param base_file_name: The batch name.
+        """
+        super().__init__(parent_dir_path, 
+                         base_file_name, 
+                         _BatchExtension.FITS)   
+        
+        
+    def _read(
+        self
+    ) -> Spectrogram:
+        """Parses a FITS file to generate a `Spectrogram` instance.
+
+        Reverses the spectra along the frequency axis and converts units to linearised 
+        values if necessary. Infers the spectrum type from the `BUNIT` header.
+
+        :raises NotImplementedError: If the `BUNIT` header value represents an unsupported spectrum type.
+        :return: A `Spectrogram` instance containing the parsed FITS file data.
+        """
+        with fits.open(self.file_path, mode='readonly') as hdulist:
+            primary_hdu                = self._get_primary_hdu(hdulist)
+            dynamic_spectra            = self._get_dynamic_spectra(primary_hdu)
+            spectrogram_start_datetime = self._get_spectrogram_start_datetime(primary_hdu)
+            bintable_hdu               = self._get_bintable_hdu(hdulist)
+            times                      = self._get_times(bintable_hdu)
+            frequencies                = self._get_frequencies(bintable_hdu)
+            bunit                      = self._get_bunit(primary_hdu)
+
+            # bunit is interpreted as a SpectrumUnit
+            spectrum_unit = SpectrumUnit(bunit)
+            if spectrum_unit == SpectrumUnit.DIGITS:
+                dynamic_spectra_linearised = self._convert_units_to_linearised(dynamic_spectra)
+                
+                return Spectrogram(dynamic_spectra_linearised[::-1, :], # reverse the spectra along the frequency axis
+                                   times, 
+                                   frequencies[::-1], # sort the frequencies in ascending order
+                                   self.tag, 
+                                   spectrum_unit,
+                                   spectrogram_start_datetime)
+            else:
+                raise NotImplementedError(f"SPECTRE does not currently support spectrum type with BUNITS '{spectrum_unit}'")   
+
+
+    def _get_primary_hdu(
+        self, hdulist: HDUList
+    ) -> PrimaryHDU:
+        return hdulist['PRIMARY']
+    
+
+    def _get_bintable_hdu(
+        self, 
+        hdulist: HDUList
+    ) -> BinTableHDU:
+        return hdulist[1]
+    
+    
+    def _get_dynamic_spectra(
+        self, 
+        primary_hdu: PrimaryHDU
+    ) -> npt.NDArray[np.float32]:
+        return primary_hdu.data.astype(np.float32)
+
+
+    def _get_spectrogram_start_datetime(
+        self, 
+        primary_hdu: PrimaryHDU
+    ) -> datetime:
+        date_obs = primary_hdu.header['DATE-OBS']
+        time_obs = primary_hdu.header['TIME-OBS']
+        return datetime.strptime(f"{date_obs}T{time_obs}", "%Y/%m/%dT%H:%M:%S.%f")
+
+
+    def _get_bunit(
+        self, 
+        primary_hdu: PrimaryHDU
+    ) -> str:
+        return primary_hdu.header['BUNIT']
+
+
+    def _convert_units_to_linearised(
+        self, 
+        raw_digits: npt.NDArray[np.float32]
+    ) -> npt.NDArray[np.float32]:
+        """Converts spectrogram data from raw digit values to linearised units.
+
+        Applies a transformation based on ADC specifications to convert raw values 
+        to dB and then to linearised units.
+
+        :param dynamic_spectra: Raw dynamic spectra in digit values.
+        :return: The dynamic spectra with linearised units.
+        """
+        # conversion as per ADC specs [see email from C. Monstein]
+        dB = (raw_digits / 255) * (2500 / 25)
+        return 10 ** (dB / 10)
+    
+    
+    def _get_times(
+        self, 
+        bintable_hdu: BinTableHDU
+    ) -> npt.NDArray[np.float32]:
+        """Extracts the elapsed times for each spectrum in seconds, with the first spectrum set to t=0
+        by convention.
+        """
+        return bintable_hdu.data['TIME'][0] # already in seconds
+
+
+    def _get_frequencies(
+        self,
+        bintable_hdu: BinTableHDU
+    ) -> npt.NDArray[np.float32]:
+        """Extracts the frequencies for each spectral component."""
+        frequencies_MHz = bintable_hdu.data['FREQUENCY'][0]
+        return frequencies_MHz * 1e6 # convert to Hz
+    
+  
+@register_batch(BatchKey.CALLISTO)
+class CallistoBatch(BaseBatch):
+    """A batch of data generated by the e-Callisto network.
+    
+    Supports the following file extensions:
+    - `.fits` (via the `spectrogram_file` attribute)
+    """
+    def __init__(
+        self,
+        start_time: str,
+        tag: str
+    ) -> None:
+        """Initialise a `CallistoBatch` instance.
+
+        :param start_time: The start time of the batch.
+        :param tag: The batch name tag.
+        """
+        super().__init__(start_time, tag) 
+        self._fits_file = _FitsFile(self.parent_dir_path, self.name)
+        
+        # add files formally to the batch
+        self.add_file( self.spectrogram_file )
+    
+    
+    @property
+    def spectrogram_file(
+        self
+    ) -> _FitsFile:
+        """The batch file corresponding to the `.fits` extension."""
+        return self._fits_file