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

spectre_core/batches/_batches.py

@@ -10,22 +10,23 @@ 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
-from spectre_core.exceptions import (
-    BatchNotFoundError
-)
-from ._base import BaseBatch, parse_batch_base_file_name
+from spectre_core.exceptions import BatchNotFoundError
+from ._base import BaseBatch, parse_batch_file_name
+
+T = TypeVar("T", bound=BaseBatch)
+
 
-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, 
+        self,
         tag: str,
         batch_cls: Type[T],
-        year: Optional[int] = None, 
-        month: Optional[int] = None, 
-        day: Optional[int] = None
+        year: Optional[int] = None,
+        month: Optional[int] = None,
+        day: Optional[int] = None,
     ) -> None:
         """Initialise a `Batches` instance.
 
@@ -40,85 +41,54 @@ class Batches(Generic[T]):
         self._batch_map: dict[str, T] = OrderedDict()
         self.set_date(year, month, day)
 
-
     @property
-    def tag(
-        self
-    ) -> str:
+    def tag(self) -> str:
         """The batch name tag."""
         return self._tag
 
-
     @property
-    def batch_cls(
-        self
-    ) -> Type[T]:
+    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]:
+    def year(self) -> Optional[int]:
         """The numeric year, to filter batch files."""
         return self._year
 
-
-    @property 
-    def month(
-        self
-    ) -> Optional[int]:
+    @property
+    def month(self) -> Optional[int]:
         """The numeric month of the year, to filter batch files."""
         return self._month
-    
 
     @property
-    def day(
-        self
-    ) -> Optional[int]:
+    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:
+    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[T]:
+    def batch_list(self) -> list[T]:
         """A list of all batches found within `batches_dir_path`."""
-        return  list(self._batch_map.values())
-    
+        return list(self._batch_map.values())
 
     @property
-    def start_times(
-        self
-    ) -> list[str]:
+    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:
+    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]
+        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.
@@ -132,72 +102,57 @@ class Batches(Generic[T]):
         self._day = day
         self.update()
 
-
-    def update(
-        self
-    ) -> None:
-        """Perform a fresh search all files in `batches_dir_path` for batches 
+    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() 
-        
+        self._batch_map = OrderedDict()
+
         # get a list of all batch file names in the batches directory path
-        batch_file_names = [f for (_, _, files) in os.walk(self.batches_dir_path) for f in files]
+        batch_file_names = [
+            f for (_, _, files) in os.walk(self.batches_dir_path) for f in files
+        ]
         for batch_file_name in batch_file_names:
-            start_time, tag, _ = parse_batch_base_file_name(batch_file_name) 
+            start_time, tag, _ = parse_batch_file_name(batch_file_name)
             if tag == self._tag:
                 self._batch_map[start_time] = self.batch_cls(start_time, tag)
-        
+
         self._batch_map = OrderedDict(sorted(self._batch_map.items()))
-    
 
-    def __iter__(
-        self
-    ) -> Iterator[T]:
+    def __iter__(self) -> Iterator[T]:
         """Iterate over the stored batch instances."""
         yield from self.batch_list
-        
-    
-    def __len__(
-        self
-    ):
-        return self.num_batches
 
+    def __len__(self):
+        return self.num_batches
 
-    def _get_from_start_time(
-        self, 
-        start_time: str
-    ) -> T:
+    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}")
+            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
-    ) -> T:
+    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")
         elif index > self.num_batches:
-            raise IndexError(f"Index '{index}' is greater than the number of batches '{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
-    ) -> T:
+    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. 
+        :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.
         """
@@ -205,22 +160,19 @@ class Batches(Generic[T]):
             return self._get_from_start_time(subscript)
         elif isinstance(subscript, int):
             return self._get_from_index(subscript)
-   
 
     def get_spectrogram(
-        self,
-        start_datetime: datetime, 
-        end_datetime: datetime
+        self, start_datetime: datetime, end_datetime: datetime
     ) -> Spectrogram:
         """
-        Retrieve a spectrogram spanning the specified time range. 
+        Retrieve a spectrogram spanning the specified time range.
 
         :param start_datetime: The start time of the range (inclusive).
         :param end_datetime: 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.
         """
-        
+
         spectrograms = []
         for batch in self:
             # skip batches without spectrogram data
@@ -233,10 +185,14 @@ class Batches(Generic[T]):
 
             # Check if the batch overlaps with the input time range
             if start_datetime <= upper_bound and lower_bound <= end_datetime:
-                spectrograms.append( time_chop(spectrogram, start_datetime, end_datetime) )
+                spectrograms.append(
+                    time_chop(spectrogram, start_datetime, end_datetime)
+                )
 
         if spectrograms:
             return join_spectrograms(spectrograms)
         else:
-            raise FileNotFoundError(f"No spectrogram data found for the time range "
-                                    f"{start_datetime} to {end_datetime}.")
+            raise FileNotFoundError(
+                f"No spectrogram data found for the time range "
+                f"{start_datetime} to {end_datetime}."
+            )

spectre_core/batches/_factory.py

@@ -11,25 +11,22 @@ 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]:
-    ...
+) -> Type[CallistoBatch]: ...
 
 
 @overload
 def get_batch_cls(
     batch_key: Literal[BatchKey.IQ_STREAM],
-) -> Type[IQStreamBatch]:
-    ...
-    
-    
+) -> Type[IQStreamBatch]: ...
+
+
 @overload
-def get_batch_cls(batch_key: BatchKey) -> Type[BaseBatch]:
-    ...
+def get_batch_cls(batch_key: BatchKey) -> Type[BaseBatch]: ...
 
 
 def get_batch_cls(
@@ -44,14 +41,14 @@ def get_batch_cls(
     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}. "
-                                 f"Valid batch keys are: {valid_batch_keys}")
+        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
-) -> Type[BaseBatch]:
+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:
@@ -61,9 +58,13 @@ def get_batch_cls_from_tag(
     else:
         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 )
+            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

@@ -10,21 +10,21 @@ 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]]:
+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]:
+
+    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

@@ -4,13 +4,14 @@
 
 from enum import Enum
 
-class BatchKey(Enum):  
+
+class BatchKey(Enum):
     """Key bound to a `Batch` plugin class.
-    
-    :ivar IQ_STREAM: Represents the default batch data generated  by `spectre`, 
+
+    :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"
-    
+
+    IQ_STREAM = "iq_stream"
+    CALLISTO = "callisto"

spectre_core/batches/plugins/_callisto.py

@@ -21,107 +21,89 @@ 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:
+
+    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:
+        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 
+        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)
+        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)
+                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}'")   
+                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_primary_hdu(
-        self, hdulist: HDUList
-    ) -> PrimaryHDU:
-        return hdulist['PRIMARY']
-    
-
-    def _get_bintable_hdu(
-        self, 
-        hdulist: HDUList
-    ) -> BinTableHDU:
+    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_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']
+    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 _get_bunit(self, primary_hdu: PrimaryHDU) -> str:
+        return primary_hdu.header["BUNIT"]
 
     def _convert_units_to_linearised(
-        self, 
-        raw_digits: npt.NDArray[np.float32]
+        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 
+        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.
@@ -130,54 +112,40 @@ class _FitsFile(BatchFile[Spectrogram]):
         # 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]:
+
+    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
+        return bintable_hdu.data["TIME"][0]  # already in seconds
 
-
-    def _get_frequencies(
-        self,
-        bintable_hdu: BinTableHDU
-    ) -> npt.NDArray[np.float32]:
+    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
-    
-  
+        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:
+
+    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) 
+        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 )
-    
-    
+        self.add_file(self.spectrogram_file)
+
     @property
-    def spectrogram_file(
-        self
-    ) -> _FitsFile:
+    def spectrogram_file(self) -> _FitsFile:
         """The batch file corresponding to the `.fits` extension."""
         return self._fits_file