ml4gw 0.7.7__py3-none-any.whl → 0.7.8__py3-none-any.whl

ml4gw/spectral.py

@@ -1,4 +1,7 @@
 """
+This module provides functions for calculation of spectral densities
+and for whitening.
+
 Several implementation details are derived from the scipy csd and welch
 implementations. For more info, see
 
@@ -9,8 +12,6 @@ and
 https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.csd.html
 """
 
-from typing import Optional, Union
-
 import torch
 from jaxtyping import Float
 from torch import Tensor
@@ -36,13 +37,11 @@ def median(x: Float[Tensor, "... size"], axis: int) -> Float[Tensor, "..."]:
     return torch.quantile(x, q=0.5, axis=axis) / bias
 
 
-def _validate_shapes(
-    x: Tensor, nperseg: int, y: Optional[Tensor] = None
-) -> None:
+def _validate_shapes(x: Tensor, nperseg: int, y: Tensor | None = None) -> None:
     if x.shape[-1] < nperseg:
         raise ValueError(
-            "Number of samples {} in input x is insufficient "
-            "for number of fft samples {}".format(x.shape[-1], nperseg)
+            f"Number of samples {x.shape[-1]} in input x is insufficient "
+            f"for number of fft samples {nperseg}"
         )
     elif x.ndim > 3:
         raise ValueError(
@@ -59,30 +58,30 @@ def _validate_shapes(
     if x.shape[-1] != y.shape[-1]:
         raise ValueError(
             "Time dimensions of x and y tensors must "
-            "be the same, found {} and {}".format(x.shape[-1], y.shape[-1])
+            f"be the same, found {x.shape[-1]} and {y.shape[-1]}"
         )
     elif x.ndim == 1 and not y.ndim == 1:
         raise ValueError(
             "Can't compute cross spectral density of "
-            "1D tensor x with {}D tensor y".format(y.ndim)
+            f"1D tensor x with {y.ndim}D tensor y"
         )
     elif x.ndim > 1 and y.ndim == x.ndim:
         if not y.shape == x.shape:
             raise ValueError(
                 "If x and y tensors have the same number "
                 "of dimensions, shapes must fully match. "
-                "Found shapes {} and {}".format(x.shape, y.shape)
+                f"Found shapes {x.shape} and {y.shape}"
             )
     elif x.ndim > 1 and y.ndim != (x.ndim - 1):
         raise ValueError(
             "Can't compute cross spectral density of "
-            "tensors with shapes {} and {}".format(x.shape, y.shape)
+            f"tensors with shapes {x.shape} and {y.shape}"
         )
     elif x.ndim > 2 and y.shape[0] != x.shape[0]:
         raise ValueError(
             "If x is a 3D tensor and y is a 2D tensor, "
             "0th batch dimensions must match, but found "
-            "values {} and {}".format(x.shape[0], y.shape[0])
+            f"values {x.shape[0]} and {y.shape[0]}"
         )
 
 
@@ -93,7 +92,7 @@ def fast_spectral_density(
     window: Float[Tensor, " {nperseg//2+1}"],
     scale: float,
     average: str = "median",
-    y: Optional[TimeSeries1to3d] = None,
+    y: TimeSeries1to3d | None = None,
 ) -> FrequencySeries1to3d:
     """
     Compute the power spectral density of a multichannel
@@ -340,10 +339,10 @@ def spectral_density(
 
 def truncate_inverse_power_spectrum(
     psd: PSDTensor,
-    fduration: Union[Float[Tensor, " time"], float],
+    fduration: Float[Tensor, " time"] | float,
     sample_rate: float,
-    highpass: Optional[float] = None,
-    lowpass: Optional[float] = None,
+    highpass: float | None = None,
+    lowpass: float | None = None,
 ) -> PSDTensor:
     """
     Truncate the length of the time domain response
@@ -460,10 +459,10 @@ def normalize_by_psd(
 def whiten(
     X: WaveformTensor,
     psd: PSDTensor,
-    fduration: Union[Float[Tensor, " time"], float],
+    fduration: Float[Tensor, " time"] | float,
     sample_rate: float,
-    highpass: Optional[float] = None,
-    lowpass: Optional[float] = None,
+    highpass: float | None = None,
+    lowpass: float | None = None,
 ) -> WaveformTensor:
     """
     Whiten a batch of timeseries using the specified
@@ -522,10 +521,8 @@ def whiten(
     N = X.size(-1)
     if N <= (2 * pad):
         raise ValueError(
-            (
-                "Not enough timeseries samples {} for number of "
-                "padded samples {}"
-            ).format(N, 2 * pad)
+            f"Not enough timeseries samples {N} for number of "
+            f"padded samples {2 * pad}"
         )
 
     # normalize the number of expected dimensions in the PSD

ml4gw/transforms/__init__.py

@@ -1,3 +1,9 @@
+"""
+This module contains a variety of data transformation classes,
+including objects to calculate spectral densities, whiten data,
+and compute Q-transforms.
+"""
+
 from .iirfilter import IIRFilter
 from .pearson import ShiftedPearsonCorrelation
 from .qtransform import QScan, SingleQTransform

ml4gw/transforms/decimator.py

@@ -0,0 +1,183 @@
+import torch
+
+
+class Decimator(torch.nn.Module):
+    r"""
+    Downsample (decimate) a timeseries according to a user-defined schedule.
+
+    .. note::
+
+        This is a naive decimator that does not use any IIR/FIR filtering
+        and selects every M-th sample according to the schedule.
+
+    The schedule specifies which segments of the input to keep and at what
+    sampling rate. Each row of the schedule has the form:
+
+        `[start_time, end_time, target_sample_rate]`
+
+    Args:
+        sample_rate (int):
+            Sampling rate (Hz) of the input timeseries.
+        schedule (torch.Tensor):
+            Tensor of shape `(N, 3)` defining start time, end time,
+            and target sample rate for each segment.
+
+    Shape:
+        - Input: `(B, C, T)` where
+            - B = batch size
+            - C = channels
+            - T = number of timesteps
+                    (must equal schedule duration × sample_rate)
+        - Output:
+            - If ``split=False`` → `(B, C, T')` where `T'` is total
+              number of decimated samples across all segments.
+            - If ``split=True`` → list of tensors, one per segment.
+
+    Returns:
+        torch.Tensor or List[torch.Tensor]:
+            The decimated timeseries, or list of decimated segments if
+            ``split=True``.
+
+    Example:
+        .. code-block:: python
+
+            >>> import torch
+            >>> from ml4gw.transforms.decimator import Decimator
+
+            >>> sample_rate = 2048
+            >>> X_duration = 60
+
+            >>> schedule = torch.tensor(
+            ...     [[0, 40, 256], [40, 58, 512], [58, 60, 2048]],
+            ...     dtype=torch.int,
+            ... )
+
+            >>> decimator = Decimator(sample_rate=sample_rate,
+            ...    schedule=schedule)
+            >>> X = torch.randn(1, 1, sample_rate * X_duration)
+            >>> X_dec = decimator(X)
+            >>> X_seg = decimator(X, split=True)
+
+            >>> print("Original shape:", X.shape)
+            Original shape: torch.Size([1, 1, 122880])
+            >>> print("Decimated shape:", X_dec.shape)
+            Decimated shape: torch.Size([1, 1, 23552])
+            >>> for i, seg in enumerate(X_seg):
+            ...     print(f"Segment {i} shape:", seg.shape)
+            Segment 0 shape: torch.Size([1, 1, 10240])
+            Segment 1 shape: torch.Size([1, 1, 9216])
+            Segment 2 shape: torch.Size([1, 1, 4096])
+    """
+
+    def __init__(
+        self,
+        sample_rate: int = None,
+        schedule: torch.Tensor = None,
+    ) -> None:
+        super().__init__()
+        self.sample_rate = sample_rate
+        self.schedule = schedule
+
+        self._validate_inputs()
+        idx = self.build_variable_indices()
+        self.register_buffer("idx", idx)
+
+        self.expected_len = int(
+            (self.schedule[:, 1][-1] - self.schedule[:, 0][0])
+            * self.sample_rate
+        )
+
+    def _validate_inputs(self) -> None:
+        r"""
+        Validate the schedule and sample_rate.
+        """
+        if self.schedule.ndim != 2 or self.schedule.shape[1] != 3:
+            raise ValueError(
+                f"Schedule must be of shape (N, 3), got {self.schedule.shape}"
+            )
+
+        if not torch.all(self.schedule[:, 1] > self.schedule[:, 0]):
+            raise ValueError(
+                "Each schedule segment must have end_time > start_time"
+            )
+
+        if torch.any(self.sample_rate % self.schedule[:, 2].long() != 0):
+            raise ValueError(
+                f"Sample rate {self.sample_rate} must be divisible by all "
+                f"target rates {self.schedule[:, 2].tolist()}"
+            )
+
+    def build_variable_indices(self) -> torch.Tensor:
+        r"""
+        Compute the time indices to keep based on the schedule.
+
+        Returns:
+            torch.Tensor:
+                1D tensor of indices used to decimate the input.
+        """
+        idx = torch.tensor([], dtype=torch.long)
+
+        for s in self.schedule:
+            if idx.size(0) == 0:
+                start = int(s[0] * self.sample_rate)
+            else:
+                start = int(idx[-1]) + int(idx[-1] - idx[-2])
+            stop = int(start + (s[1] - s[0]) * self.sample_rate)
+            step = int(self.sample_rate // s[2])
+            new_idx = torch.arange(start, stop, step, dtype=torch.long)
+            idx = torch.cat((idx, new_idx))
+        return idx
+
+    def split_by_schedule(self, X: torch.Tensor) -> tuple[torch.Tensor, ...]:
+        r"""
+        Split a decimated timeseries into segments according to the schedule.
+
+        Args:
+            X (torch.Tensor):
+                Decimated input of shape `(B, C, T')`.
+
+        Returns:
+            tuple of torch.Tensor:
+                Each segment has shape :math:`(B, C, T_i)`
+                where :math:`T_i` is the length implied by
+                the corresponding schedule row.
+        """
+        split_sizes = (
+            ((self.schedule[:, 1] - self.schedule[:, 0]) * self.schedule[:, 2])
+            .long()
+            .tolist()
+        )
+
+        return torch.split(X, split_sizes, dim=-1)
+
+    def forward(
+        self,
+        X: torch.Tensor,
+        split: bool = False,
+    ) -> torch.Tensor | list[torch.Tensor]:
+        r"""
+        Apply decimation to the input timeseries.
+
+        Args:
+            X (torch.Tensor):
+                Input tensor of shape `(B, C, T)`, where `T` must equal
+                schedule duration × sample_rate.
+            split (bool, optional):
+                If True, return a list of segments instead of a single
+                concatenated tensor. Default: False.
+
+        Returns:
+            torch.Tensor or List[torch.Tensor]:
+                Decimated timeseries, or list of decimated segments.
+        """
+        if X.shape[-1] != self.expected_len:
+            raise ValueError(
+                f"X length {X.shape[-1]} does not match "
+                f"expected schedule duration {self.expected_len}"
+            )
+
+        X_dec = X.index_select(dim=-1, index=self.idx)
+
+        if split:
+            X_dec = self.split_by_schedule(X_dec)
+        return X_dec

ml4gw/transforms/iirfilter.py

@@ -1,5 +1,3 @@
-from typing import Union
-
 import torch
 from scipy.signal import iirfilter
 from torchaudio.functional import filtfilt
@@ -55,9 +53,9 @@ class IIRFilter(torch.nn.Module):
     def __init__(
         self,
         N: int,
-        Wn: Union[float, torch.Tensor],
-        rs: Union[None, float, torch.Tensor] = None,
-        rp: Union[None, float, torch.Tensor] = None,
+        Wn: float | torch.Tensor,
+        rs: None | float | torch.Tensor = None,
+        rp: None | float | torch.Tensor = None,
         btype="band",
         analog=False,
         ftype="butter",

ml4gw/transforms/pearson.py

@@ -52,15 +52,14 @@ class ShiftedPearsonCorrelation(torch.nn.Module):
             raise ValueError(
                 "y may not have more dimensions that x for "
                 "ShiftedPearsonCorrelation, but found shapes "
-                "{} and {}".format(y.shape, x.shape)
+                f"{y.shape} and {x.shape}"
             )
         for dim in range(y.ndim):
             if y.size(-dim - 1) != x.size(-dim - 1):
                 raise ValueError(
                     "x and y expected to have same size along "
-                    "last dimensions, but found shapes {} and {}".format(
-                        x.shape, y.shape
-                    )
+                    f"last dimensions, but found shapes {x.shape} and "
+                    f"{y.shape}"
                 )
 
     def forward(

ml4gw/transforms/qtransform.py

@@ -1,6 +1,5 @@
 import math
 import warnings
-from typing import List, Tuple
 
 import torch
 import torch.nn.functional as F
@@ -146,7 +145,7 @@ class QTile(torch.nn.Module):
                 means = torch.mean(energy, dim=-1, keepdim=True)
                 energy /= means
             else:
-                raise ValueError("Invalid normalisation %r" % norm)
+                raise ValueError(f"Invalid normalisation {norm}")
             energy = energy.type(torch.float32)
         return energy
 
@@ -194,9 +193,9 @@ class SingleQTransform(torch.nn.Module):
         self,
         duration: float,
         sample_rate: float,
-        spectrogram_shape: Tuple[int, int],
+        spectrogram_shape: tuple[int, int],
         q: float = 12,
-        frange: List[float] = None,
+        frange: list[float] = None,
         mismatch: float = 0.2,
         interpolation_method: str = "bicubic",
     ) -> None:
@@ -303,7 +302,7 @@ class SingleQTransform(torch.nn.Module):
         return torch.unique(freqs)
 
     def get_max_energy(
-        self, fsearch_range: List[float] = None, dimension: str = "both"
+        self, fsearch_range: list[float] = None, dimension: str = "both"
     ):
         """
         Gets the maximum energy value among the QTiles. The maximum can
@@ -375,7 +374,7 @@ class SingleQTransform(torch.nn.Module):
                 [
                     qtile_interpolator(qtile)
                     for qtile, qtile_interpolator in zip(
-                        qtiles, self.qtile_interpolators
+                        qtiles, self.qtile_interpolators, strict=True
                     )
                 ],
                 dim=-2,
@@ -459,9 +458,9 @@ class QScan(torch.nn.Module):
         self,
         duration: float,
         sample_rate: float,
-        spectrogram_shape: Tuple[int, int],
-        qrange: List[float] = None,
-        frange: List[float] = None,
+        spectrogram_shape: tuple[int, int],
+        qrange: list[float] = None,
+        frange: list[float] = None,
         interpolation_method="bicubic",
         mismatch: float = 0.2,
     ) -> None:
@@ -500,7 +499,7 @@ class QScan(torch.nn.Module):
             ]
         )
 
-    def get_qs(self) -> List[float]:
+    def get_qs(self) -> list[float]:
         """
         Determine the values of Q to try for the set of Q-transforms
         """
@@ -518,7 +517,7 @@ class QScan(torch.nn.Module):
     def forward(
         self,
         X: TimeSeries1to3d,
-        fsearch_range: List[float] = None,
+        fsearch_range: list[float] = None,
         norm: str = "median",
     ):
         """

ml4gw/transforms/scaler.py

@@ -1,5 +1,3 @@
-from typing import Optional
-
 import torch
 from jaxtyping import Float
 from torch import Tensor
@@ -24,7 +22,7 @@ class ChannelWiseScaler(FittableTransform):
             to be 1D (single channel).
     """
 
-    def __init__(self, num_channels: Optional[int] = None) -> None:
+    def __init__(self, num_channels: int | None = None) -> None:
         super().__init__()
 
         shape = (num_channels or 1,)
@@ -37,7 +35,7 @@ class ChannelWiseScaler(FittableTransform):
         self.register_buffer("std", std)
 
     def fit(
-        self, X: Float[Tensor, "... time"], std_reg: Optional[float] = 0.0
+        self, X: Float[Tensor, "... time"], std_reg: float | None = 0.0
     ) -> None:
         """Fit the scaling parameters to a timeseries
 
@@ -59,7 +57,7 @@ class ChannelWiseScaler(FittableTransform):
         else:
             raise ValueError(
                 "Can't fit channel wise mean and standard deviation "
-                "from tensor of shape {}".format(X.shape)
+                f"from tensor of shape {X.shape}"
             )
         std += std_reg * torch.ones_like(std)
         super().build(mean=mean, std=std)

ml4gw/transforms/snr_rescaler.py

@@ -1,5 +1,3 @@
-from typing import Optional
-
 import torch
 
 from ..gw import compute_network_snr
@@ -13,8 +11,8 @@ class SnrRescaler(FittableSpectralTransform):
         num_channels: int,
         sample_rate: float,
         waveform_duration: float,
-        highpass: Optional[float] = None,
-        lowpass: Optional[float] = None,
+        highpass: float | None = None,
+        lowpass: float | None = None,
         dtype: torch.dtype = torch.float32,
     ) -> None:
         super().__init__()
@@ -45,15 +43,13 @@ class SnrRescaler(FittableSpectralTransform):
     def fit(
         self,
         *background: TimeSeries2d,
-        fftlength: Optional[float] = None,
-        overlap: Optional[float] = None,
+        fftlength: float | None = None,
+        overlap: float | None = None,
     ):
         if len(background) != self.num_channels:
             raise ValueError(
-                "Expected to fit whitening transform on {} background "
-                "timeseries, but was passed {}".format(
-                    self.num_channels, len(background)
-                )
+                f"Expected to fit whitening transform on {self.num_channels} "
+                f"background timeseries, but was passed {len(background)}"
             )
 
         num_freqs = self.background.size(1)
@@ -69,7 +65,7 @@ class SnrRescaler(FittableSpectralTransform):
     def forward(
         self,
         responses: WaveformTensor,
-        target_snrs: Optional[BatchTensor] = None,
+        target_snrs: BatchTensor | None = None,
     ):
         snrs = compute_network_snr(
             responses,

ml4gw/transforms/spectral.py

@@ -1,5 +1,3 @@
-from typing import Optional
-
 import torch
 from jaxtyping import Float
 from torch import Tensor
@@ -52,20 +50,17 @@ class SpectralDensity(torch.nn.Module):
         self,
         sample_rate: float,
         fftlength: float,
-        overlap: Optional[float] = None,
+        overlap: float | None = None,
         average: str = "mean",
-        window: Optional[
-            Float[Tensor, " {int(fftlength*sample_rate)}"]
-        ] = None,
+        window: Float[Tensor, " {int(fftlength*sample_rate)}"] | None = None,
         fast: bool = False,
     ) -> None:
         if overlap is None:
             overlap = fftlength / 2
         elif overlap >= fftlength:
             raise ValueError(
-                "Can't have overlap {} longer than fftlength {}".format(
-                    overlap, fftlength
-                )
+                f"Can't have overlap {overlap} longer than fftlength "
+                f"{fftlength}"
             )
 
         super().__init__()
@@ -80,9 +75,7 @@ class SpectralDensity(torch.nn.Module):
 
         if window.size(0) != self.nperseg:
             raise ValueError(
-                "Window must have length {} got {}".format(
-                    self.nperseg, window.size(0)
-                )
+                f"Window must have length {self.nperseg} got {window.size(0)}"
             )
         self.register_buffer("window", window)
 
@@ -99,7 +92,7 @@ class SpectralDensity(torch.nn.Module):
         self.fast = fast
 
     def forward(
-        self, x: TimeSeries1to3d, y: Optional[TimeSeries1to3d] = None
+        self, x: TimeSeries1to3d, y: TimeSeries1to3d | None = None
     ) -> FrequencySeries1to3d:
         if self.fast:
             return fast_spectral_density(

ml4gw/transforms/spectrogram.py

@@ -1,5 +1,4 @@
 import warnings
-from typing import Dict, List
 
 import torch
 import torch.nn.functional as F
@@ -104,7 +103,7 @@ class MultiResolutionSpectrogram(torch.nn.Module):
         self.register_buffer("freq_idxs", freq_idxs)
         self.register_buffer("time_idxs", time_idxs)
 
-    def _check_and_format_kwargs(self, kwargs: Dict[str, List]) -> List:
+    def _check_and_format_kwargs(self, kwargs: dict[str, list]) -> list:
         lengths = sorted(len(v) for v in kwargs.values())
         lengths = list(set(lengths))
 
@@ -127,7 +126,10 @@ class MultiResolutionSpectrogram(torch.nn.Module):
             size = lengths[1]
             kwargs = {k: v * int(size / len(v)) for k, v in kwargs.items()}
 
-        return [dict(zip(kwargs, col)) for col in zip(*kwargs.values())]
+        return [
+            dict(zip(kwargs, col, strict=True))
+            for col in zip(*kwargs.values(), strict=True)
+        ]
 
     def forward(
         self, X: TimeSeries3d
@@ -161,6 +163,7 @@ class MultiResolutionSpectrogram(torch.nn.Module):
             self.right_pad,
             self.top_pad,
             self.bottom_pad,
+            strict=True,
         ):
             padded_specs.append(F.pad(spec, (left, right, top, bottom)))
 

ml4gw/transforms/spline_interpolation.py

@@ -2,8 +2,6 @@
 Adaptation of code from https://github.com/dottormale/Qtransform_torch/
 """
 
-from typing import Optional, Tuple
-
 import torch
 from torch import Tensor
 
@@ -50,7 +48,7 @@ class SplineInterpolateBase(torch.nn.Module):
         t: Tensor,
         d: int,
         m: int,
-    ) -> Tuple[Tensor, Tensor]:
+    ) -> tuple[Tensor, Tensor]:
         """
         Compute the L and R values for B-spline basis functions.
         L and R are respectively the first and second coefficient multiplying
@@ -208,7 +206,7 @@ class SplineInterpolate1D(SplineInterpolateBase):
         x_in: Tensor,
         kx: int = 3,
         sx: float = 0.0,
-        x_out: Optional[Tensor] = None,
+        x_out: Tensor | None = None,
     ):
         super().__init__()
 
@@ -284,7 +282,7 @@ class SplineInterpolate1D(SplineInterpolateBase):
     def forward(
         self,
         Z: Tensor,
-        x_out: Optional[Tensor] = None,
+        x_out: Tensor | None = None,
     ) -> Tensor:
         """
         Compute the interpolated data
@@ -377,8 +375,8 @@ class SplineInterpolate2D(SplineInterpolateBase):
         ky: int = 3,
         sx: float = 0.0,
         sy: float = 0.0,
-        x_out: Optional[Tensor] = None,
-        y_out: Optional[Tensor] = None,
+        x_out: Tensor | None = None,
+        y_out: Tensor | None = None,
     ):
         super().__init__()
 
@@ -483,8 +481,8 @@ class SplineInterpolate2D(SplineInterpolateBase):
     def forward(
         self,
         Z: Tensor,
-        x_out: Optional[Tensor] = None,
-        y_out: Optional[Tensor] = None,
+        x_out: Tensor | None = None,
+        y_out: Tensor | None = None,
     ) -> Tensor:
         """
         Compute the interpolated data

ml4gw/transforms/transform.py

@@ -1,5 +1,3 @@
-from typing import Optional
-
 import torch
 
 from ..spectral import spectral_density
@@ -20,8 +18,8 @@ class FittableTransform(torch.nn.Module):
     def _check_built(self):
         if not self.built:
             raise ValueError(
-                "Must fit parameters of {} transform to data "
-                "before calling forward step".format(self.__class__.__name__)
+                f"Must fit parameters of {self.__class__.__name__} transform "
+                "to data before calling forward step"
             )
 
     def __call__(self, *args, **kwargs):
@@ -47,8 +45,8 @@ class FittableSpectralTransform(FittableTransform):
         x: TimeSeries1to3d,
         sample_rate: float,
         num_freqs: int,
-        fftlength: Optional[float] = None,
-        overlap: Optional[float] = None,
+        fftlength: float | None = None,
+        overlap: float | None = None,
     ) -> FrequencySeries1to3d:
         # if we specified an FFT length, convert
         # the (assumed) time-domain data to the