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

spectre_core/batches/_batches.py

@@ -0,0 +1,247 @@
+# SPDX-FileCopyrightText: © 2024 Jimmy Fitzpatrick <jcfitzpatrick12@gmail.com>
+# This file is part of SPECTRE
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+import os
+from typing import Optional, TypeVar, Type, Generic, Iterator
+from collections import OrderedDict
+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
+
+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_cls = batch_cls
+        self._batch_map: dict[str, T] = OrderedDict()
+        self.set_date(year, month, day)
+
+
+    @property
+    def tag(
+        self
+    ) -> str:
+        """The batch name tag."""
+        return self._tag
+
+
+    @property
+    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
+    ) -> Optional[int]:
+        """The numeric month of the year, to filter batch files."""
+        return self._month
+    
+
+    @property
+    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 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]:
+        """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 found within `batches_dir_path`."""
+        return list(self._batch_map.keys())
+
+
+    @property
+    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:
+        """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()
+
+
+    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() 
+        
+        # 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]
+        for batch_file_name in batch_file_names:
+            # strip the extension
+            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_cls(start_time, tag)
+        
+        self._batch_map = OrderedDict(sorted(self._batch_map.items()))
+    
+
+    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
+    ) -> 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
+    ) -> 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}'")
+        return self.batch_list[index]
+
+
+    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 get_spectrogram_from_range(
+        self,
+        start_time: str, 
+        end_time: str
+    ) -> Spectrogram:
+        """
+        Retrieve a spectrogram spanning the specified 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, TimeFormat.DATETIME)
+        end_datetime   = datetime.strptime(end_time,   TimeFormat.DATETIME)
+
+        spectrograms = []
+        for batch in self:
+            # skip batches without spectrogram data
+            if not batch.spectrogram_file.exists:
+                continue
+
+            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:
+                spectrograms.append( time_chop(spectrogram, start_time, end_time) )
+
+        if spectrograms:
+            return join_spectrograms(spectrograms)
+        else:
+            raise FileNotFoundError(f"No spectrogram data found for the time range "
+                                    f"{start_time} to {end_time}.")

spectre_core/batches/_factory.py

@@ -0,0 +1,69 @@
+# 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 Literal, overload, Type, cast
+
+from spectre_core.exceptions import BatchNotFoundError
+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]:
+    ...
+
+
+@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}. "
+                                 f"Valid batch keys are: {valid_batch_keys}")
+    return batch_cls
+
+
+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:
+        return get_batch_cls(BatchKey.CALLISTO)
+
+    # otherwise, assume that the tag has an associated capture config,
+    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 )

spectre_core/batches/_register.py

@@ -0,0 +1,30 @@
+# 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 Type, Callable, TypeVar
+
+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

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