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

ml4gw/nn/streaming/snapshotter.py

@@ -1,4 +1,4 @@
-from typing import Optional, Sequence, Tuple
+from collections.abc import Sequence
 
 import torch
 from jaxtyping import Float
@@ -58,15 +58,13 @@ class Snapshotter(torch.nn.Module):
         snapshot_size: int,
         stride_size: int,
         batch_size: int,
-        channels_per_snapshot: Optional[Sequence[int]] = None,
+        channels_per_snapshot: Sequence[int] | None = None,
     ) -> None:
         super().__init__()
         if stride_size >= snapshot_size:
             raise ValueError(
-                "Snapshotter can't accommodate stride {} "
-                "which is greater than snapshot size {}".format(
-                    stride_size, snapshot_size
-                )
+                f"Snapshotter can't accommodate stride {stride_size} "
+                f"which is greater than snapshot size {snapshot_size}"
             )
 
         self.snapshot_size = snapshot_size
@@ -77,9 +75,8 @@ class Snapshotter(torch.nn.Module):
         if channels_per_snapshot is not None:
             if sum(channels_per_snapshot) != num_channels:
                 raise ValueError(
-                    "Can't break {} channels into {}".format(
-                        num_channels, channels_per_snapshot
-                    )
+                    f"Can't break {num_channels} channels into "
+                    f"{channels_per_snapshot}"
                 )
         self.channels_per_snapshot = channels_per_snapshot
         self.num_channels = num_channels
@@ -90,8 +87,8 @@ class Snapshotter(torch.nn.Module):
     def forward(
         self,
         update: Float[Tensor, "channel time1"],
-        snapshot: Optional[Float[Tensor, "channel time2"]] = None,
-    ) -> Tuple[Tensor, ...]:
+        snapshot: Float[Tensor, "channel time2"] | None = None,
+    ) -> tuple[Tensor, ...]:
         if snapshot is None:
             snapshot = self.get_initial_state()
 
@@ -108,9 +105,8 @@ class Snapshotter(torch.nn.Module):
         if self.channels_per_snapshot is not None:
             if snapshots.size(1) != self.num_channels:
                 raise ValueError(
-                    "Expected {} channels, found {}".format(
-                        self.num_channels, snapshots.size(1)
-                    )
+                    f"Expected {self.num_channels} channels, found "
+                    f"{snapshots.size(1)}"
                 )
             snapshots = torch.split(
                 snapshots, self.channels_per_snapshot, dim=1

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
@@ -5,6 +11,6 @@ from .scaler import ChannelWiseScaler
 from .snr_rescaler import SnrRescaler
 from .spectral import SpectralDensity
 from .spectrogram import MultiResolutionSpectrogram
-from .spline_interpolation import SplineInterpolate
+from .spline_interpolation import SplineInterpolate1D, SplineInterpolate2D
 from .waveforms import WaveformProjector, WaveformSampler
 from .whitening import FixedWhiten, Whiten

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
@@ -8,7 +7,7 @@ from jaxtyping import Float, Int
 from torch import Tensor
 
 from ..types import FrequencySeries1to3d, TimeSeries1to3d, TimeSeries3d
-from .spline_interpolation import SplineInterpolate
+from .spline_interpolation import SplineInterpolate1D
 
 """
 All based on https://github.com/gwpy/gwpy/blob/v3.0.8/gwpy/signal/qtransform.py
@@ -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:
@@ -260,18 +259,15 @@ class SingleQTransform(torch.nn.Module):
         )
         self.qtile_interpolators = torch.nn.ModuleList(
             [
-                SplineInterpolate(
+                SplineInterpolate1D(
                     kx=3,
                     x_in=torch.arange(0, self.duration, self.duration / tiles),
-                    y_in=torch.arange(len(idx)),
                     x_out=t_out,
-                    y_out=torch.arange(len(idx)),
                 )
-                for tiles, idx in zip(unique_ntiles, self.stack_idx)
+                for tiles in unique_ntiles
             ]
         )
 
-        t_in = t_out
         f_in = self.freqs
         f_out = torch.logspace(
             math.log10(self.frange[0]),
@@ -279,13 +275,10 @@ class SingleQTransform(torch.nn.Module):
             self.spectrogram_shape[0],
         )
 
-        self.interpolator = SplineInterpolate(
+        self.interpolator = SplineInterpolate1D(
             kx=3,
-            ky=3,
-            x_in=t_in,
-            y_in=f_in,
-            x_out=t_out,
-            y_out=f_out,
+            x_in=f_in,
+            x_out=f_out,
         )
 
     def get_freqs(self) -> Float[Tensor, " nfreq"]:
@@ -309,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
@@ -379,14 +372,15 @@ class SingleQTransform(torch.nn.Module):
             ]
             time_interped = torch.cat(
                 [
-                    interpolator(qtile)
-                    for qtile, interpolator in zip(
-                        qtiles, self.qtile_interpolators
+                    qtile_interpolator(qtile)
+                    for qtile, qtile_interpolator in zip(
+                        qtiles, self.qtile_interpolators, strict=True
                     )
                 ],
                 dim=-2,
             )
-            return self.interpolator(time_interped)
+            # Transpose because the final dimension gets interpolated
+            return self.interpolator(time_interped.mT).mT
         num_f_bins, num_t_bins = self.spectrogram_shape
         resampled = [
             F.interpolate(
@@ -464,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:
@@ -505,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
         """
@@ -523,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(