ml4gw 0.5.0__py3-none-any.whl → 0.5.1__py3-none-any.whl

ml4gw/transforms/pearson.py

@@ -1,5 +1,8 @@
 import torch
+from jaxtyping import Float
+from torch import Tensor
 
+from ml4gw.types import TimeSeries1to3d
 from ml4gw.utils.slicing import unfold_windows
 
 
@@ -40,7 +43,7 @@ class ShiftedPearsonCorrelation(torch.nn.Module):
         super().__init__()
         self.max_shift = max_shift
 
-    def _shape_checks(self, x: torch.Tensor, y: torch.Tensor):
+    def _shape_checks(self, x: TimeSeries1to3d, y: TimeSeries1to3d):
         if x.ndim > 3:
             raise ValueError(
                 "Tensor x can only have up to 3 dimensions "
@@ -61,8 +64,9 @@ class ShiftedPearsonCorrelation(torch.nn.Module):
                     )
                 )
 
-    # TODO: torchtyping annotate
-    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+    def forward(
+        self, x: TimeSeries1to3d, y: TimeSeries1to3d
+    ) -> Float[Tensor, "windows ..."]:
         self._shape_checks(x, y)
         dim = x.size(-1)
 

ml4gw/transforms/qtransform.py

@@ -3,6 +3,10 @@ from typing import List, Optional, Tuple
 
 import torch
 import torch.nn.functional as F
+from jaxtyping import Float, Int
+from torch import Tensor
+
+from ml4gw.types import FrequencySeries1to3d, TimeSeries1to3d, TimeSeries3d
 
 """
 All based on https://github.com/gwpy/gwpy/blob/v3.0.8/gwpy/signal/qtransform.py
@@ -44,7 +48,7 @@ class QTile(torch.nn.Module):
         duration: float,
         sample_rate: float,
         mismatch: float,
-    ):
+    ) -> None:
         super().__init__()
         self.mismatch = mismatch
         self.q = q
@@ -63,18 +67,18 @@ class QTile(torch.nn.Module):
         self.register_buffer("indices", self.get_data_indices())
         self.register_buffer("window", self.get_window())
 
-    def ntiles(self):
+    def ntiles(self) -> int:
         """
         Number of tiles in this frequency row
         """
         tcum_mismatch = self.duration * 2 * torch.pi * self.frequency / self.q
         return int(2 ** torch.ceil(torch.log2(tcum_mismatch / self.deltam)))
 
-    def _get_indices(self):
+    def _get_indices(self) -> Int[Tensor, " windowsize"]:
         half = int((self.windowsize - 1) / 2)
         return torch.arange(-half, half + 1)
 
-    def get_window(self):
+    def get_window(self) -> Float[Tensor, " windowsize"]:
         """
         Generate the bi-square window for this row
         """
@@ -87,7 +91,7 @@ class QTile(torch.nn.Module):
         )
         return torch.Tensor((1 - xfrequencies**2) ** 2 * norm)
 
-    def get_data_indices(self):
+    def get_data_indices(self) -> Int[Tensor, " windowsize"]:
         """
         Get the index array of relevant frequencies for this row
         """
@@ -95,7 +99,9 @@ class QTile(torch.nn.Module):
             self._get_indices() + 1 + self.frequency * self.duration,
         ).type(torch.long)
 
-    def forward(self, fseries: torch.Tensor, norm: str = "median"):
+    def forward(
+        self, fseries: FrequencySeries1to3d, norm: str = "median"
+    ) -> TimeSeries1to3d:
         """
         Compute the transform for this row
 
@@ -176,7 +182,7 @@ class SingleQTransform(torch.nn.Module):
         q: float = 12,
         frange: List[float] = [0, torch.inf],
         mismatch: float = 0.2,
-    ):
+    ) -> None:
         super().__init__()
         self.q = q
         self.spectrogram_shape = spectrogram_shape
@@ -198,7 +204,7 @@ class SingleQTransform(torch.nn.Module):
         )
         self.qtiles = None
 
-    def get_freqs(self):
+    def get_freqs(self) -> Float[Tensor, " nfreq"]:
         """
         Calculate the frequencies that will be used in this transform.
         For each frequency, a `QTile` is created.
@@ -262,7 +268,7 @@ class SingleQTransform(torch.nn.Module):
         if dimension == "batch":
             return torch.max(max_across_ft, dim=-1).values
 
-    def compute_qtiles(self, X: torch.Tensor, norm: str = "median"):
+    def compute_qtiles(self, X: TimeSeries1to3d, norm: str = "median") -> None:
         """
         Take the FFT of the input timeseries and calculate the transform
         for each `QTile`
@@ -272,7 +278,7 @@ class SingleQTransform(torch.nn.Module):
         X[..., 1:] *= 2
         self.qtiles = [qtile(X, norm) for qtile in self.qtile_transforms]
 
-    def interpolate(self, num_f_bins: int, num_t_bins: int):
+    def interpolate(self, num_f_bins: int, num_t_bins: int) -> TimeSeries3d:
         """
         Interpolate each `QTile` to the specified number of time and
         frequency bins. Note that PyTorch does not have the same
@@ -299,7 +305,7 @@ class SingleQTransform(torch.nn.Module):
 
     def forward(
         self,
-        X: torch.Tensor,
+        X: TimeSeries1to3d,
         norm: str = "median",
         spectrogram_shape: Optional[Tuple[int, int]] = None,
     ):
@@ -371,7 +377,7 @@ class QScan(torch.nn.Module):
         qrange: List[float] = [4, 64],
         frange: List[float] = [0, torch.inf],
         mismatch: float = 0.2,
-    ):
+    ) -> None:
         super().__init__()
         self.qrange = qrange
         self.mismatch = mismatch
@@ -397,7 +403,7 @@ class QScan(torch.nn.Module):
             ]
         )
 
-    def get_qs(self):
+    def get_qs(self) -> List[float]:
         """
         Determine the values of Q to try for the set of Q-transforms
         """
@@ -413,7 +419,7 @@ class QScan(torch.nn.Module):
 
     def forward(
         self,
-        X: torch.Tensor,
+        X: TimeSeries1to3d,
         fsearch_range: List[float] = None,
         norm: str = "median",
         spectrogram_shape: Optional[Tuple[int, int]] = None,

ml4gw/transforms/scaler.py

@@ -1,6 +1,8 @@
 from typing import Optional
 
 import torch
+from jaxtyping import Float
+from torch import Tensor
 
 from ml4gw.transforms.transform import FittableTransform
 
@@ -34,7 +36,7 @@ class ChannelWiseScaler(FittableTransform):
         self.register_buffer("mean", mean)
         self.register_buffer("std", std)
 
-    def fit(self, X: torch.Tensor) -> None:
+    def fit(self, X: Float[Tensor, "... time"]) -> None:
         """Fit the scaling parameters to a timeseries
 
         Computes the channel-wise mean and standard deviation
@@ -60,7 +62,9 @@ class ChannelWiseScaler(FittableTransform):
 
         super().build(mean=mean, std=std)
 
-    def forward(self, X: torch.Tensor, reverse: bool = False) -> torch.Tensor:
+    def forward(
+        self, X: Float[Tensor, "... time"], reverse: bool = False
+    ) -> Float[Tensor, "... time"]:
         if not reverse:
             return (X - self.mean) / self.std
         else:

ml4gw/transforms/snr_rescaler.py

@@ -2,8 +2,9 @@ from typing import Optional
 
 import torch
 
-from ml4gw import gw
+from ml4gw.gw import compute_network_snr
 from ml4gw.transforms.transform import FittableSpectralTransform
+from ml4gw.types import BatchTensor, TimeSeries2d, WaveformTensor
 
 
 class SnrRescaler(FittableSpectralTransform):
@@ -34,7 +35,7 @@ class SnrRescaler(FittableSpectralTransform):
 
     def fit(
         self,
-        *background: torch.Tensor,
+        *background: TimeSeries2d,
         fftlength: Optional[float] = None,
         overlap: Optional[float] = None,
     ):
@@ -58,10 +59,10 @@ class SnrRescaler(FittableSpectralTransform):
 
     def forward(
         self,
-        responses: gw.WaveformTensor,
-        target_snrs: Optional[gw.ScalarTensor] = None,
+        responses: WaveformTensor,
+        target_snrs: Optional[BatchTensor] = None,
     ):
-        snrs = gw.compute_network_snr(
+        snrs = compute_network_snr(
             responses, self.background, self.sample_rate, self.mask
         )
         if target_snrs is None:

ml4gw/transforms/spectral.py

@@ -1,8 +1,11 @@
 from typing import Optional
 
 import torch
+from jaxtyping import Float
+from torch import Tensor
 
 from ml4gw.spectral import fast_spectral_density, spectral_density
+from ml4gw.types import FrequencySeries1to3d, TimeSeries1to3d
 
 
 class SpectralDensity(torch.nn.Module):
@@ -51,7 +54,9 @@ class SpectralDensity(torch.nn.Module):
         fftlength: float,
         overlap: Optional[float] = None,
         average: str = "mean",
-        window: Optional[torch.Tensor] = None,
+        window: Optional[
+            Float[Tensor, " {int(fftlength*sample_rate)}"]
+        ] = None,
         fast: bool = False,
     ) -> None:
         if overlap is None:
@@ -93,7 +98,9 @@ class SpectralDensity(torch.nn.Module):
         self.average = average
         self.fast = fast
 
-    def forward(self, x: torch.Tensor, y: Optional[torch.Tensor] = None):
+    def forward(
+        self, x: TimeSeries1to3d, y: Optional[TimeSeries1to3d] = None
+    ) -> FrequencySeries1to3d:
         if self.fast:
             return fast_spectral_density(
                 x,

ml4gw/transforms/spectrogram.py

@@ -3,8 +3,12 @@ from typing import Dict, List
 
 import torch
 import torch.nn.functional as F
+from jaxtyping import Float
+from torch import Tensor
 from torchaudio.transforms import Spectrogram
 
+from ml4gw.types import TimeSeries3d
+
 
 class MultiResolutionSpectrogram(torch.nn.Module):
     """
@@ -122,7 +126,9 @@ class MultiResolutionSpectrogram(torch.nn.Module):
 
         return [dict(zip(kwargs, col)) for col in zip(*kwargs.values())]
 
-    def forward(self, X: torch.Tensor) -> torch.Tensor:
+    def forward(
+        self, X: TimeSeries3d
+    ) -> Float[Tensor, "batch channel frequency time"]:
         """
         Calculate spectrograms of the input tensor and
         combine them into a single spectrogram

ml4gw/transforms/transform.py

@@ -3,6 +3,7 @@ from typing import Optional
 import torch
 
 from ml4gw.spectral import spectral_density
+from ml4gw.types import FrequencySeries1to3d, TimeSeries1to3d
 
 
 class FittableTransform(torch.nn.Module):
@@ -43,12 +44,12 @@ class FittableTransform(torch.nn.Module):
 class FittableSpectralTransform(FittableTransform):
     def normalize_psd(
         self,
-        x,
+        x: TimeSeries1to3d,
         sample_rate: float,
         num_freqs: int,
         fftlength: Optional[float] = None,
         overlap: Optional[float] = None,
-    ):
+    ) -> FrequencySeries1to3d:
         # if we specified an FFT length, convert
         # the (assumed) time-domain data to the
         # frequency domain
@@ -68,7 +69,7 @@ class FittableSpectralTransform(FittableTransform):
                 scale=scale,
             )
 
-        # add two dummy dimensions in case we need to inerpolate
+        # add two dummy dimensions in case we need to interpolate
         # the frequency dimension, since `interpolate` expects
         # a (batch, channel, spatial) formatted tensor as input
         x = x.view(1, 1, -1)

ml4gw/transforms/waveforms.py

@@ -1,8 +1,11 @@
 from typing import List, Optional
 
 import torch
+from jaxtyping import Float
+from torch import Tensor
 
 from ml4gw import gw
+from ml4gw.types import BatchTensor
 
 
 # TODO: should these live in ml4gw.waveforms submodule?
@@ -10,8 +13,8 @@ from ml4gw import gw
 class WaveformSampler(torch.nn.Module):
     def __init__(
         self,
-        parameters: Optional[torch.Tensor] = None,
-        **polarizations: torch.Tensor,
+        parameters: Optional[Float[Tensor, "batch num_params"]] = None,
+        **polarizations: Float[Tensor, "batch time"],
     ):
         super().__init__()
         # make sure we have the same number of waveforms
@@ -29,7 +32,7 @@ class WaveformSampler(torch.nn.Module):
             elif num_waveforms is None:
                 num_waveforms = tensor.shape[0]
 
-            self.polarizations[polarization] = torch.Tensor(tensor)
+            self.polarizations[polarization] = Tensor(tensor)
 
         if parameters is not None and len(parameters) != num_waveforms:
             raise ValueError(
@@ -73,10 +76,10 @@ class WaveformProjector(torch.nn.Module):
 
     def forward(
         self,
-        dec: gw.ScalarTensor,
-        psi: gw.ScalarTensor,
-        phi: gw.ScalarTensor,
-        **polarizations,
+        dec: BatchTensor,
+        psi: BatchTensor,
+        phi: BatchTensor,
+        **polarizations: Float[Tensor, "batch time"],
     ):
         ifo_responses = gw.compute_observed_strain(
             dec,

ml4gw/transforms/whitening.py

@@ -1,9 +1,15 @@
-from typing import Optional
+from typing import Optional, Union
 
 import torch
 
 from ml4gw import spectral
 from ml4gw.transforms.transform import FittableSpectralTransform
+from ml4gw.types import (
+    FrequencySeries1d,
+    FrequencySeries1to3d,
+    TimeSeries1d,
+    TimeSeries3d,
+)
 
 
 class Whiten(torch.nn.Module):
@@ -58,7 +64,9 @@ class Whiten(torch.nn.Module):
         window = torch.hann_window(size, dtype=torch.float64)
         self.register_buffer("window", window)
 
-    def forward(self, X: torch.Tensor, psd: torch.Tensor) -> torch.Tensor:
+    def forward(
+        self, X: TimeSeries3d, psd: FrequencySeries1to3d
+    ) -> TimeSeries3d:
         """
         Whiten a batch of multichannel timeseries by a
         background power spectral density.
@@ -142,7 +150,7 @@ class FixedWhiten(FittableSpectralTransform):
     def fit(
         self,
         fduration: float,
-        *background: torch.Tensor,
+        *background: Union[TimeSeries1d, FrequencySeries1d],
         fftlength: Optional[float] = None,
         highpass: Optional[float] = None,
         overlap: Optional[float] = None
@@ -224,7 +232,7 @@ class FixedWhiten(FittableSpectralTransform):
         fduration = torch.Tensor([fduration])
         self.build(psd=psd, fduration=fduration)
 
-    def forward(self, X: torch.Tensor) -> torch.Tensor:
+    def forward(self, X: TimeSeries3d) -> TimeSeries3d:
         """
         Whiten the input timeseries tensor using the
         PSD fit by the `.fit` method, which must be

ml4gw/types.py

@@ -1,10 +1,25 @@
-from torchtyping import TensorType
-
-WaveformTensor = TensorType["batch", "num_ifos", "time"]
-PSDTensor = TensorType["num_ifos", "frequency"]
-ScalarTensor = TensorType["batch"]
-VectorGeometry = TensorType["batch", "space"]
-TensorGeometry = TensorType["batch", "space", "space"]
-NetworkVertices = TensorType["num_ifos", 3]
-NetworkDetectorTensors = TensorType["num_ifos", 3, 3]
-TimeSeriesTensor = TensorType["num_channels", "time"]
+from typing import Union
+
+from jaxtyping import Float
+from torch import Tensor
+
+WaveformTensor = Float[Tensor, "batch num_ifos time"]
+PSDTensor = Float[Tensor, "num_ifos frequency"]
+BatchTensor = Float[Tensor, "batch"]
+VectorGeometry = Float[Tensor, "batch space"]
+TensorGeometry = Float[Tensor, "batch space space"]
+NetworkVertices = Float[Tensor, "num_ifos 3"]
+NetworkDetectorTensors = Float[Tensor, "num_ifos 3 3"]
+
+
+TimeSeries1d = Float[Tensor, "time"]
+TimeSeries2d = Float[TimeSeries1d, "channel"]
+TimeSeries3d = Float[TimeSeries2d, "batch"]
+TimeSeries1to3d = Union[TimeSeries1d, TimeSeries2d, TimeSeries3d]
+
+FrequencySeries1d = Float[Tensor, "frequency"]
+FrequencySeries2d = Float[FrequencySeries1d, "channel"]
+FrequencySeries3d = Float[FrequencySeries2d, "batch"]
+FrequencySeries1to3d = Union[
+    FrequencySeries1d, FrequencySeries2d, FrequencySeries3d
+]

ml4gw/utils/interferometer.py

@@ -4,7 +4,7 @@ import torch
 # based on values from
 # https://lscsoft.docs.ligo.org/lalsuite/lal/_l_a_l_detectors_8h_source.html
 class InterferometerGeometry:
-    def __init__(self, name: str):
+    def __init__(self, name: str) -> None:
         if name == "H1":
             self.x_arm = torch.Tensor(
                 (-0.22389266154, +0.79983062746, +0.55690487831)

ml4gw/utils/slicing.py

@@ -1,25 +1,30 @@
 from typing import Optional, Union
 
 import torch
+from jaxtyping import Float, Int64
+from torch import Tensor
 from torch.nn.functional import unfold
-from torchtyping import TensorType
 
-# need to define these for flake8 compatibility
-batch = time = channel = None  # noqa
+from ml4gw.types import (
+    TimeSeries1d,
+    TimeSeries1to3d,
+    TimeSeries2d,
+    TimeSeries3d,
+)
 
-TimeSeriesTensor = Union[TensorType["time"], TensorType["channel", "time"]]
-
-BatchTimeSeriesTensor = Union[
-    TensorType["batch", "time"], TensorType["batch", "channel", "time"]
-]
+BatchTimeSeriesTensor = Union[Float[Tensor, "batch time"], TimeSeries3d]
 
 
 def unfold_windows(
-    x: torch.Tensor,
+    x: TimeSeries1to3d,
     window_size: int,
     stride: int,
     drop_last: bool = True,
-):
+) -> Union[
+    Float[TimeSeries1d, " window"],
+    Float[TimeSeries2d, " window"],
+    Float[TimeSeries3d, " window"],
+]:
     """Unfold a timeseries into windows
 
     Args:
@@ -83,8 +88,8 @@ def unfold_windows(
 
 
 def slice_kernels(
-    x: Union[TimeSeriesTensor, TensorType["batch", "channel", "time"]],
-    idx: TensorType[..., torch.int64],
+    x: TimeSeries1to3d,
+    idx: Int64[Tensor, "..."],
     kernel_size: int,
 ) -> BatchTimeSeriesTensor:
     """Slice kernels from single or multichannel timeseries
@@ -96,7 +101,8 @@ def slice_kernels(
     one more dimension than `x`.
 
     Args:
-        x: The timeseries tensor to slice kernels from
+        x:
+            The timeseries tensor to slice kernels from
         idx:
             The indices in `x` of the first sample of each
             kernel. If `x` is 1D, `idx` must be 1D as well.
@@ -114,6 +120,7 @@ def slice_kernels(
             coincidentally among the channels.
         kernel_size:
             The length of the kernels to slice from the timeseries
+
     Returns:
         A tensor of shape `(batch_size, kernel_size)` if `x` is
         1D and `(batch_size, num_channels, kernel_size)` if `x`
@@ -225,7 +232,7 @@ def slice_kernels(
 
 
 def sample_kernels(
-    X: TimeSeriesTensor,
+    X: TimeSeries1to3d,
     kernel_size: int,
     N: Optional[int] = None,
     max_center_offset: Optional[int] = None,
@@ -245,8 +252,9 @@ def sample_kernels(
     either be `None` or be equal to `len(X)`.
 
     Args:
-        X: The timeseries tensor from which to sample kernels
-        kernel_size: The size of the kernels to sample
+        X:
+            The timeseries tensor from which to sample kernels
+            kernel_size: The size of the kernels to sample
         N:
             The number of kernels to sample. Can be left as
             `None` if `X` is 3D, otherwise must be specified

ml4gw/waveforms/generator.py

@@ -1,24 +1,26 @@
-from typing import Callable
+from typing import Callable, Dict, Tuple
 
 import torch
+from jaxtyping import Float
+from torch import Tensor
 
 
 class ParameterSampler(torch.nn.Module):
-    def __init__(self, **parameters: Callable):
+    def __init__(self, **parameters: Callable) -> None:
         super().__init__()
         self.parameters = parameters
 
     def forward(
         self,
         N: int,
-    ):
+    ) -> Dict[str, Float[Tensor, " {N}"]]:
         return {k: v.sample((N,)) for k, v in self.parameters.items()}
 
 
 class WaveformGenerator(torch.nn.Module):
     def __init__(
         self, waveform: Callable, parameter_sampler: ParameterSampler
-    ):
+    ) -> None:
         """
         A torch module that generates waveforms from a given waveform function
         and a parameter sampler.
@@ -34,6 +36,8 @@ class WaveformGenerator(torch.nn.Module):
         self.waveform = waveform
         self.parameter_sampler = parameter_sampler
 
-    def forward(self, N: int):
+    def forward(
+        self, N: int
+    ) -> Tuple[Float[Tensor, "{N} samples"], Dict[str, Float[Tensor, " {N}"]]]:
         parameters = self.parameter_sampler(N)
         return self.waveform(**parameters), parameters

ml4gw/waveforms/phenom_d.py

@@ -1,7 +1,9 @@
 import torch
-from torchtyping import TensorType
+from jaxtyping import Float
+
+from ml4gw.constants import MTSUN_SI, PI
+from ml4gw.types import BatchTensor, FrequencySeries1d
 
-from ..constants import MTSUN_SI, PI
 from .phenom_d_data import QNMData_a, QNMData_fdamp, QNMData_fring
 from .taylorf2 import TaylorF2
 
@@ -15,14 +17,14 @@ class IMRPhenomD(TaylorF2):
 
     def forward(
         self,
-        f: TensorType,
-        chirp_mass: TensorType,
-        mass_ratio: TensorType,
-        chi1: TensorType,
-        chi2: TensorType,
-        distance: TensorType,
-        phic: TensorType,
-        inclination: TensorType,
+        f: FrequencySeries1d,
+        chirp_mass: BatchTensor,
+        mass_ratio: BatchTensor,
+        chi1: BatchTensor,
+        chi2: BatchTensor,
+        distance: BatchTensor,
+        phic: BatchTensor,
+        inclination: BatchTensor,
         f_ref: float,
     ):
         """
@@ -76,15 +78,15 @@ class IMRPhenomD(TaylorF2):
 
     def phenom_d_htilde(
         self,
-        f: TensorType,
-        chirp_mass: TensorType,
-        mass_ratio: TensorType,
-        chi1: TensorType,
-        chi2: TensorType,
-        distance: TensorType,
-        phic: TensorType,
+        f: FrequencySeries1d,
+        chirp_mass: BatchTensor,
+        mass_ratio: BatchTensor,
+        chi1: BatchTensor,
+        chi2: BatchTensor,
+        distance: BatchTensor,
+        phic: BatchTensor,
         f_ref: float,
-    ):
+    ) -> Float[FrequencySeries1d, " batch"]:
         total_mass = chirp_mass * (1 + mass_ratio) ** 1.2 / mass_ratio**0.6
         mass_1 = total_mass / (1 + mass_ratio)
         mass_2 = mass_1 * mass_ratio