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

spectre_core/spectrograms/_transform.py

@@ -3,28 +3,38 @@
 # SPDX-License-Identifier: GPL-3.0-or-later
 
 import numpy as np
+import numpy.typing as npt
 from datetime import datetime, timedelta
 from typing import Optional
 from math import floor
 
-from spectre_core.config import TimeFormats
-from ._array_operations import find_closest_index, average_array
+from spectre_core.config import TimeFormat
+from ._array_operations import find_closest_index, average_array, time_elapsed
 from ._spectrogram import Spectrogram
 
-
-def frequency_chop(input_spectrogram: Spectrogram, 
-                   start_frequency: float | int, 
-                   end_frequency: float | int) -> Optional[Spectrogram]:
-    
-    is_entirely_below_frequency_range = (start_frequency < input_spectrogram.frequencies[0] and end_frequency < input_spectrogram.frequencies[0])
-    is_entirely_above_frequency_range = (start_frequency > input_spectrogram.frequencies[-1] and end_frequency > input_spectrogram.frequencies[-1])
-    # if the requested frequency range is out of bounds for the spectrogram return None
+def frequency_chop(
+    spectrogram: Spectrogram, 
+    start_frequency  : float, 
+    end_frequency    : float
+) -> Spectrogram:
+    """
+    Extracts a portion of the spectrogram within the specified frequency range.
+
+    :param spectrogram: The input spectrogram to process.
+    :param start_frequency: The starting frequency of the desired range (Hz).
+    :param end_frequency: The ending frequency of the desired range (Hz).
+    :raises ValueError: If the specified frequency range is entirely outside the spectrogram's frequency range.
+    :raises ValueError: If the start and end indices for the frequency range are identical.
+    :return: A new spectrogram containing only the specified frequency range.
+    """
+    is_entirely_below_frequency_range = (start_frequency < spectrogram.frequencies[0] and end_frequency < spectrogram.frequencies[0])
+    is_entirely_above_frequency_range = (start_frequency > spectrogram.frequencies[-1] and end_frequency > spectrogram.frequencies[-1])
     if is_entirely_below_frequency_range or is_entirely_above_frequency_range:
-        return None
+        raise ValueError(f"The requested frequency interval is entirely out of range of the input spectrogram.")
     
-    #find the index of the nearest matching frequency bins in the spectrogram
-    start_index = find_closest_index(start_frequency, input_spectrogram.frequencies)
-    end_index = find_closest_index(end_frequency, input_spectrogram.frequencies)
+    # find the index of the nearest matching frequency bins in the spectrogram
+    start_index = find_closest_index(np.float32(start_frequency), spectrogram.frequencies)
+    end_index   = find_closest_index(np.float32(end_frequency)  , spectrogram.frequencies)
     
     # enforce distinct start and end indices
     if start_index == end_index:
@@ -35,163 +45,193 @@ def frequency_chop(input_spectrogram: Spectrogram,
         start_index, end_index = end_index, start_index
     
     # chop the spectrogram accordingly
-    transformed_dynamic_spectra = input_spectrogram.dynamic_spectra[start_index:end_index+1, :]
-    transformed_frequencies = input_spectrogram.frequencies[start_index:end_index+1]
+    transformed_dynamic_spectra = spectrogram.dynamic_spectra[start_index:end_index+1, :]
+    transformed_frequencies     = spectrogram.frequencies[start_index:end_index+1]
     
-    # return the spectrogram instance
     return Spectrogram(transformed_dynamic_spectra,
-                       input_spectrogram.times,
+                       spectrogram.times,
                        transformed_frequencies,
-                       input_spectrogram.tag,
-                       input_spectrogram.start_datetime,
-                       input_spectrogram.spectrum_type)
-
-
-def time_chop(input_spectrogram: Spectrogram, 
-              start_time: str, 
-              end_time: str, 
-              time_format: str = TimeFormats.DATETIME) -> Optional[Spectrogram]:
-    
-    # parse the strings as datetimes
-    start_datetime = datetime.strptime(start_time, time_format)
-    end_datetime = datetime.strptime(end_time, time_format)
-
-    # if the requested time range is out of bounds for the spectrogram return None
-    is_entirely_below_time_range = (start_datetime < input_spectrogram.datetimes[0] and end_datetime < input_spectrogram.datetimes[0])
-    is_entirely_above_time_range = (start_datetime > input_spectrogram.datetimes[-1] and end_datetime > input_spectrogram.datetimes[-1])
+                       spectrogram.tag,
+                       spectrogram.spectrum_unit,
+                       spectrogram.start_datetime)
+
+
+def time_chop(
+    spectrogram: Spectrogram, 
+    start_time: str, 
+    end_time: str
+) -> Spectrogram:
+    """
+    Extracts a portion of the spectrogram within the specified time range.
+
+    :param spectrogram: The input spectrogram to process.
+    :param start_time: The starting time of the desired range.
+    :param end_time: The ending time of the desired range.
+    :raises ValueError: If the specified time range is entirely outside the spectrogram's time range.
+    :raises ValueError: If the start and end indices for the time range are identical.
+    :return: A new spectrogram containing only the specified time range.
+    """
+    start_datetime = np.datetime64( datetime.strptime(start_time, TimeFormat.DATETIME) )
+    end_datetime   = np.datetime64( datetime.strptime(end_time  , TimeFormat.DATETIME) )
+
+    is_entirely_below_time_range = (start_datetime < spectrogram.datetimes[0] and end_datetime < spectrogram.datetimes[0])
+    is_entirely_above_time_range = (start_datetime > spectrogram.datetimes[-1] and end_datetime > spectrogram.datetimes[-1])
     if is_entirely_below_time_range or is_entirely_above_time_range:
-        return None
+        raise ValueError(f"Requested time interval is entirely out of range of the input spectrogram.")
     
-    start_index = find_closest_index(start_datetime, input_spectrogram.datetimes)
-    end_index = find_closest_index(end_datetime, input_spectrogram.datetimes)
+    # find the index of the nearest matching spectrums in the spectrogram.
+    start_index = find_closest_index(start_datetime, spectrogram.datetimes)
+    end_index   = find_closest_index(end_datetime, spectrogram.datetimes)
     
+    # enforce distinct start and end indices
     if start_index == end_index:
         raise ValueError(f"Start and end indices are equal! Got start_index: {start_index} and end_index: {end_index}")
     
+    # if start index is more than end index, swap the ordering so to enforce start_index <= end_index
     if start_index > end_index:
         start_index, end_index = end_index, start_index
 
-    # chop the spectrogram 
-    transformed_dynamic_spectra = input_spectrogram.dynamic_spectra[:, start_index:end_index+1]
-
-    # compute the new start datetime following the time chop
-    transformed_start_datetime = input_spectrogram.datetimes[start_index]
+    # chop the spectrogram accordingly
+    transformed_dynamic_spectra = spectrogram.dynamic_spectra[:, start_index:end_index+1]
+    transformed_start_datetime = spectrogram.datetimes[start_index]
 
-    # chop the times array
-    transformed_times = input_spectrogram.times[start_index:end_index+1]
-    # assign the first spectrum to t=0 [s]
+    # chop the times array and translate such that the first spectrum to t=0 [s]
+    transformed_times  = spectrogram.times[start_index:end_index+1]
     transformed_times -= transformed_times[0]
 
     return Spectrogram(transformed_dynamic_spectra, 
                        transformed_times, 
-                       input_spectrogram.frequencies, 
-                       input_spectrogram.tag, 
-                       start_time = transformed_start_datetime,
-                       spectrum_type = input_spectrogram.spectrum_type)
+                       spectrogram.frequencies, 
+                       spectrogram.tag, 
+                       spectrogram.spectrum_unit,
+                       transformed_start_datetime)
+
+
+def _validate_and_compute_average_over(
+    original_resolution: float,
+    resolution: Optional[float], 
+    average_over: int
+) -> int:
+    """
+    Validates the input parameters and computes `average_over` if `resolution` is specified.
+
+    :param resolution: The desired resolution for averaging. Mutually exclusive with `average_over`.
+    :param average_over: The number of consecutive spectrums to average over. Mutually exclusive with `resolution`.
+    :param original_resolution: The original resolution (e.g., time or frequency).
+    :raises ValueError: If neither or both `resolution` and `average_over` are specified.
+    :return: The computed or validated `average_over` value.
+    """
+    if (resolution is None) and (average_over == 1):
+        return average_over
+
+    if not (resolution is not None) ^ (average_over != 1):
+        raise ValueError("Exactly one of 'resolution' or 'average_over' must be specified.")
 
-
-def time_average(input_spectrogram: Spectrogram, 
-                 resolution: Optional[float] = None,
-                 average_over: Optional[int] = None) -> Spectrogram:
-
-    # spectre does not currently support averaging of non-datetime assigned spectrograms
-    if not input_spectrogram.start_datetime_is_set:
-        raise NotImplementedError(f"Time averaging is not yet supported for spectrograms without an assigned datetime.")
-    
-    # if nothing is specified, do nothing
-    if (resolution is None) and (average_over is None):
-        average_over = 1
-
-    if not (resolution is not None) ^ (average_over is not None):
-        raise ValueError(f"Exactly one of 'resolution' or 'average_over' "
-                         f"must be specified.")
-    
-    # if the resolution is specified, compute the appropriate number of spectrums to average over
-    # and recall the same function
     if resolution is not None:
-        average_over = max(1, floor(resolution / input_spectrogram.time_resolution))
-        return time_average(input_spectrogram, average_over=average_over)
-    
-    # No averaging is required, if we have to average over every one spectrum
-    if average_over == 1:
-        return input_spectrogram
-
-    # average the dynamic spectra array
-    transformed_dynamic_spectra = average_array(input_spectrogram.dynamic_spectra, 
-                                                average_over, 
-                                                axis=1)
-
-    # We need to assign timestamps to the averaged spectrums in the spectrograms. 
-    # The natural way to do this is to assign the i'th averaged spectrogram 
-    # to the i'th averaged time
-    transformed_times = average_array(input_spectrogram.times, average_over)
-    
-    # find the new batch start time, which we will assign to the first spectrum after averaging
-    transformed_start_datetime = input_spectrogram.datetimes[0] + timedelta(seconds = float(transformed_times[0]))
-
-    # finally, translate the averaged time seconds to begin at t=0 [s]
+        return max(1, floor(resolution / original_resolution))
+
+    else:
+        return average_over
+
+
+def time_average(
+    spectrogram: Spectrogram,
+    resolution: Optional[float] = None,
+    average_over: int = 1,
+) -> Spectrogram:
+    """
+    Performs time averaging on the spectrogram data.
+
+    :param spectrogram: The input spectrogram to process.
+    :param resolution: The desired time resolution for averaging (seconds). Mutually exclusive with `average_over`.
+    :param average_over: The number of consecutive time points to average. Mutually exclusive with `resolution`.
+    :raises NotImplementedError: If the spectrogram lacks a defined start datetime.
+    :raises ValueError: If neither or both `resolution` and `average_over` are specified.
+    :return: A new spectrogram with time-averaged data.
+    """
+    if not spectrogram.start_datetime_is_set:
+        raise NotImplementedError(
+            "Time averaging is not supported for spectrograms without an assigned start datetime."
+        )
+
+    average_over = _validate_and_compute_average_over(
+        spectrogram.time_resolution,
+        resolution, 
+        average_over
+    )
+
+    # Perform averaging
+    transformed_dynamic_spectra = average_array(
+        spectrogram.dynamic_spectra, average_over, axis=1
+    )
+    transformed_times = average_array(spectrogram.times, average_over)
+
+    # Update start datetime and adjust times to start at t=0
+    transformed_start_datetime = (
+        spectrogram.datetimes[0] + (transformed_times[0]*1e6).astype("timedelta64[us]")
+    )
     transformed_times -= transformed_times[0]
-    
-    return Spectrogram(transformed_dynamic_spectra, 
-                       transformed_times, 
-                       input_spectrogram.frequencies, 
-                       input_spectrogram.tag,
-                       transformed_start_datetime,
-                       input_spectrogram.spectrum_type)
-
-
-
-def frequency_average(input_spectrogram: Spectrogram, 
-                      resolution: Optional[float] = None,
-                      average_over: Optional[int] = None) -> Spectrogram:
-    
-    # if nothing is specified, do nothing
-    if (resolution is None) and (average_over is None):
-        average_over = 1
-
-    if not (resolution is not None) ^ (average_over is not None):
-        raise ValueError(f"Exactly one of 'resolution' or 'average_over' "
-                         f"must be specified.")
-    
-    # if the resolution is specified, compute the appropriate number of spectrums to average over
-    # and recall the same function
-    if resolution is not None:
-        average_over = max(1, floor(resolution / input_spectrogram.frequency_resolution))
-        return frequency_average(input_spectrogram, average_over=average_over)
-    
-    # No averaging is required, if we have to average over every one spectrum
-    if average_over == 1:
-        return input_spectrogram
-    
-    # We need to assign physical frequencies to the averaged spectrums in the spectrograms.
-    # is to assign the i'th averaged spectral component to the i'th averaged frequency.
-    # average the dynamic spectra array
-    transformed_dynamic_spectra = average_array(input_spectrogram.dynamic_spectra, 
-                                                 average_over, 
-                                                 axis=0)
-    transformed_frequencies = average_array(input_spectrogram.frequencies, average_over)
-
-    return Spectrogram(transformed_dynamic_spectra, 
-                       input_spectrogram.times, 
-                       transformed_frequencies, 
-                       input_spectrogram.tag,
-                       input_spectrogram.start_datetime,
-                       input_spectrogram.spectrum_type)
-
-
-def _time_elapsed(datetimes: np.ndarray) -> np.ndarray:
-    # Extract the first datetime to use as the reference point
-    base_time = datetimes[0]
-    # Calculate elapsed time in seconds for each datetime in the list
-    elapsed_time = [(dt - base_time).total_seconds() for dt in datetimes]
-    # Convert the list of seconds to a NumPy array of type float32
-    return np.array(elapsed_time, dtype=np.float32)
-
-
-# we assume that the spectrogram list is ordered chronologically
-# we assume there is no time overlap in any of the spectrograms in the list
-def join_spectrograms(spectrograms: list[Spectrogram]) -> Spectrogram:
 
+    return Spectrogram(
+        transformed_dynamic_spectra,
+        transformed_times,
+        spectrogram.frequencies,
+        spectrogram.tag,
+        spectrogram.spectrum_unit,
+        transformed_start_datetime,
+    )
+
+
+def frequency_average(
+    spectrogram: Spectrogram,
+    resolution: Optional[float] = None,
+    average_over: int = 1,
+) -> Spectrogram:
+    """
+    Performs frequency averaging on the spectrogram data.
+
+    :param spectrogram: The input spectrogram to process.
+    :param resolution: The desired frequency resolution for averaging (Hz). Mutually exclusive with `average_over`.
+    :param average_over: The number of consecutive frequency bins to average. Mutually exclusive with `resolution`.
+    :raises ValueError: If neither or both `resolution` and `average_over` are specified.
+    :return: A new spectrogram with frequency-averaged data.
+    """
+    average_over = _validate_and_compute_average_over(
+        spectrogram.frequency_resolution,
+        resolution, 
+        average_over
+    )
+
+    # Perform averaging
+    transformed_dynamic_spectra = average_array(
+        spectrogram.dynamic_spectra, average_over, axis=0
+    )
+    transformed_frequencies = average_array(spectrogram.frequencies, average_over)
+
+    return Spectrogram(
+        transformed_dynamic_spectra,
+        spectrogram.times,
+        transformed_frequencies,
+        spectrogram.tag,
+        spectrogram.spectrum_unit,
+        spectrogram.start_datetime,
+    )
+
+
+def join_spectrograms(
+    spectrograms: list[Spectrogram]
+) -> Spectrogram:
+    """
+    Joins multiple spectrograms into a single spectrogram along the time axis.
+
+    :param spectrograms: A list of spectrograms to combine.
+    :raises ValueError: If the input list is empty.
+    :raises ValueError: If spectrograms have mismatched frequency ranges.
+    :raises ValueError: If spectrograms have different tags.
+    :raises ValueError: If spectrograms have differing spectrum units.
+    :raises ValueError: If any spectrogram lacks a defined start datetime.
+    :return: A new spectrogram combining all input spectrograms along the time axis.
+    """
     # check that the length of the list is non-zero
     num_spectrograms = len(spectrograms)
     if num_spectrograms == 0:
@@ -207,31 +247,19 @@ def join_spectrograms(spectrograms: list[Spectrogram]) -> Spectrogram:
             raise ValueError(f"All spectrograms must have identical frequency ranges")
         if spectrogram.tag != reference_spectrogram.tag:
             raise ValueError(f"All tags must be equal for each spectrogram in the input list!")
-        if spectrogram.spectrum_type != reference_spectrogram.spectrum_type:
+        if spectrogram.spectrum_unit != reference_spectrogram.spectrum_unit:
             raise ValueError(f"All units must be equal for each spectrogram in the input list!")
         if not spectrogram.start_datetime_is_set:
             raise ValueError(f"All spectrograms must have their start datetime set.")
 
-    # build a list of the time array of each spectrogram in the list
-    conc_datetimes = np.concatenate([s.datetimes for s in spectrograms])
-    # find the total number of time stamps
-    num_total_time_bins = len(conc_datetimes)
-    # find the total number of frequency bins (we can safely now use the first)
-    num_total_freq_bins = len(reference_spectrogram.frequencies)
-    # create an empty numpy array to hold the joined spectrograms
-    transformed_dynamic_spectra = np.empty((num_total_freq_bins, num_total_time_bins))
-
-    start_index = 0
-    for spectrogram in spectrograms:
-        end_index = start_index + len(spectrogram.times)
-        transformed_dynamic_spectra[:, start_index:end_index] = spectrogram.dynamic_spectra
-        start_index = end_index
-
-    transformed_times = _time_elapsed(conc_datetimes)
+    # Concatenate all dynamic spectra directly along the time axis
+    transformed_dynamic_spectra = np.hstack([spectrogram.dynamic_spectra for spectrogram in spectrograms])
+    
+    transformed_times = time_elapsed( np.concatenate([s.datetimes for s in spectrograms]) )
 
     return Spectrogram(transformed_dynamic_spectra, 
                        transformed_times, 
                        reference_spectrogram.frequencies, 
                        reference_spectrogram.tag, 
-                       reference_spectrogram.start_datetime,
-                       reference_spectrogram.spectrum_type) 
+                       reference_spectrogram.spectrum_unit,
+                       reference_spectrogram.start_datetime) 

spectre_core/wgetting/__init__.py

@@ -2,8 +2,10 @@
 # This file is part of SPECTRE
 # SPDX-License-Identifier: GPL-3.0-or-later
 
-from ._callisto import download_callisto_data, CALLISTO_INSTRUMENT_CODES
+"""Download third-party spectrogram data."""
+
+from ._callisto import download_callisto_data, CallistoInstrumentCode
 
 __all__ = [
-    "download_callisto_data", "CALLISTO_INSTRUMENT_CODES"
+    "download_callisto_data", "CallistoInstrumentCode"
 ]