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

ml4gw/augmentations.py

@@ -1,4 +1,6 @@
 import torch
+from jaxtyping import Float
+from torch import Tensor
 
 
 class SignalInverter(torch.nn.Module):
@@ -16,7 +18,9 @@ class SignalInverter(torch.nn.Module):
         super().__init__()
         self.prob = prob
 
-    def forward(self, X):
+    def forward(
+        self, X: Float[Tensor, "*batch time"]
+    ) -> Float[Tensor, "*batch time"]:
         mask = torch.rand(size=X.shape[:-1]) < self.prob
         X[mask] *= -1
         return X
@@ -37,7 +41,9 @@ class SignalReverser(torch.nn.Module):
         super().__init__()
         self.prob = prob
 
-    def forward(self, X):
+    def forward(
+        self, X: Float[Tensor, "*batch time"]
+    ) -> Float[Tensor, "*batch time"]:
         mask = torch.rand(size=X.shape[:-1]) < self.prob
         X[mask] = X[mask].flip(-1)
         return X

ml4gw/dataloading/chunked_dataset.py

@@ -2,6 +2,8 @@ from collections.abc import Iterable
 
 import torch
 
+from ml4gw.types import WaveformTensor
+
 
 class ChunkedTimeSeriesDataset(torch.utils.data.IterableDataset):
     """
@@ -55,10 +57,10 @@ class ChunkedTimeSeriesDataset(torch.utils.data.IterableDataset):
         self.coincident = coincident
         self.device = device
 
-    def __len__(self):
+    def __len__(self) -> int:
         return len(self.chunk_it) * self.batches_per_chunk
 
-    def __iter__(self):
+    def __iter__(self) -> WaveformTensor:
         it = iter(self.chunk_it)
         chunk = next(it)
         num_chunks, num_channels, chunk_size = chunk.shape

ml4gw/dataloading/hdf5_dataset.py

@@ -161,7 +161,7 @@ class Hdf5TimeSeriesDataset(torch.utils.data.IterableDataset):
                     x[b, c] = f[self.channels[c]][i : i + self.kernel_size]
         return torch.Tensor(x)
 
-    def __iter__(self) -> torch.Tensor:
+    def __iter__(self) -> WaveformTensor:
         worker_info = torch.utils.data.get_worker_info()
         if worker_info is None:
             num_batches = self.batches_per_epoch

ml4gw/dataloading/in_memory_dataset.py

@@ -2,8 +2,9 @@ import itertools
 from typing import Optional, Tuple, Union
 
 import torch
+from jaxtyping import Float
+from torch import Tensor
 
-from ml4gw import types
 from ml4gw.utils.slicing import slice_kernels
 
 
@@ -76,9 +77,9 @@ class InMemoryDataset(torch.utils.data.IterableDataset):
 
     def __init__(
         self,
-        X: types.TimeSeriesTensor,
+        X: Float[Tensor, "channels time"],
         kernel_size: int,
-        y: Optional[types.ScalarTensor] = None,
+        y: Optional[Float[Tensor, " time"]] = None,
         batch_size: int = 32,
         stride: int = 1,
         batches_per_epoch: Optional[int] = None,
@@ -207,7 +208,10 @@ class InMemoryDataset(torch.utils.data.IterableDataset):
 
     def __iter__(
         self,
-    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+    ) -> Union[
+        Float[Tensor, "batch channel time"],
+        Tuple[Float[Tensor, "batch channel time"], Float[Tensor, " batch"]],
+    ]:
 
         indices = self.init_indices()
         for i in range(len(self)):

ml4gw/distributions.py

@@ -9,6 +9,8 @@ from typing import Optional
 
 import torch
 import torch.distributions as dist
+from jaxtyping import Float
+from torch import Tensor
 
 
 class Cosine(dist.Distribution):
@@ -31,11 +33,11 @@ class Cosine(dist.Distribution):
         self.high = torch.as_tensor(high)
         self.norm = 1 / (torch.sin(self.high) - torch.sin(self.low))
 
-    def rsample(self, sample_shape: torch.Size = torch.Size()) -> torch.Tensor:
+    def rsample(self, sample_shape: torch.Size = torch.Size()) -> Tensor:
         u = torch.rand(sample_shape, device=self.low.device)
         return torch.arcsin(u / self.norm + torch.sin(self.low))
 
-    def log_prob(self, value):
+    def log_prob(self, value: float) -> Float[Tensor, ""]:
         value = torch.as_tensor(value)
         inside_range = (value >= self.low) & (value <= self.high)
         return value.cos().log() * inside_range
@@ -164,7 +166,7 @@ class DeltaFunction(dist.Distribution):
         super().__init__(batch_shape, validate_args=validate_args)
         self.peak = torch.as_tensor(peak)
 
-    def rsample(self, sample_shape: torch.Size = torch.Size()) -> torch.Tensor:
+    def rsample(self, sample_shape: torch.Size = torch.Size()) -> Tensor:
         return self.peak * torch.ones(
             sample_shape, device=self.peak.device, dtype=torch.float32
         )

ml4gw/gw.py

@@ -13,27 +13,21 @@ https://github.com/lscsoft/bilby/blob/master/bilby/gw/detector/interferometer.py
 from typing import List, Tuple, Union
 
 import torch
-from torchtyping import TensorType
+from jaxtyping import Float
+from torch import Tensor
 
+from ml4gw.constants import C
 from ml4gw.types import (
+    BatchTensor,
     NetworkDetectorTensors,
     NetworkVertices,
     PSDTensor,
-    ScalarTensor,
     TensorGeometry,
     VectorGeometry,
     WaveformTensor,
 )
 from ml4gw.utils.interferometer import InterferometerGeometry
 
-SPEED_OF_LIGHT = 299792458.0  # m/s
-
-
-# define some tensor shapes we'll reuse a bit
-# up front. Need to assign these variables so
-# that static linters don't give us name errors
-batch = num_ifos = polarizations = time = frequency = space = None  # noqa
-
 
 def outer(x: VectorGeometry, y: VectorGeometry) -> TensorGeometry:
     """
@@ -62,12 +56,12 @@ polarization_funcs = {
 
 
 def compute_antenna_responses(
-    theta: ScalarTensor,
-    psi: ScalarTensor,
-    phi: ScalarTensor,
+    theta: BatchTensor,
+    psi: BatchTensor,
+    phi: BatchTensor,
     detector_tensors: NetworkDetectorTensors,
     modes: List[str],
-) -> TensorType["batch", "polarizations", "num_ifos"]:
+) -> Float[Tensor, "batch polarizations num_ifos"]:
     """
     Compute the antenna pattern factors of a batch of
     waveforms as a function of the sky parameters of
@@ -147,8 +141,8 @@ def compute_antenna_responses(
 
 def shift_responses(
     responses: WaveformTensor,
-    theta: ScalarTensor,
-    phi: ScalarTensor,
+    theta: BatchTensor,
+    phi: BatchTensor,
     vertices: NetworkVertices,
     sample_rate: float,
 ) -> WaveformTensor:
@@ -166,7 +160,7 @@ def shift_responses(
     # Divide by c in the second line so that we only
     # need to multiply the array by a single float
     dt = -(omega * vertices).sum(axis=-1)
-    dt *= sample_rate / SPEED_OF_LIGHT
+    dt *= sample_rate / C
     dt = torch.trunc(dt).type(torch.int64)
 
     # rolling by gathering implementation based on
@@ -191,13 +185,13 @@ def shift_responses(
 
 
 def compute_observed_strain(
-    dec: ScalarTensor,
-    psi: ScalarTensor,
-    phi: ScalarTensor,
+    dec: BatchTensor,
+    psi: BatchTensor,
+    phi: BatchTensor,
     detector_tensors: NetworkDetectorTensors,
     detector_vertices: NetworkVertices,
     sample_rate: float,
-    **polarizations: TensorType["batch", "time"],
+    **polarizations: Float[Tensor, "batch time"],
 ) -> WaveformTensor:
     """
     Compute the strain timeseries $h(t)$ observed by a network
@@ -289,8 +283,8 @@ def compute_ifo_snr(
     responses: WaveformTensor,
     psd: PSDTensor,
     sample_rate: float,
-    highpass: Union[float, TensorType["frequency"], None] = None,
-) -> TensorType["batch", "num_ifos"]:
+    highpass: Union[float, Float[Tensor, " frequency"], None] = None,
+) -> Float[Tensor, "batch num_ifos"]:
     r"""Compute the SNRs of a batch of interferometer responses
 
     Compute the signal to noise ratio (SNR) of individual
@@ -390,8 +384,8 @@ def compute_network_snr(
     responses: WaveformTensor,
     psd: PSDTensor,
     sample_rate: float,
-    highpass: Union[float, TensorType["frequency"], None] = None,
-) -> ScalarTensor:
+    highpass: Union[float, Float[Tensor, " frequency"], None] = None,
+) -> BatchTensor:
     r"""
     Compute the total SNR from a gravitational waveform
     from a network of interferometers. The total SNR for
@@ -437,10 +431,10 @@ def compute_network_snr(
 
 def reweight_snrs(
     responses: WaveformTensor,
-    target_snrs: Union[float, ScalarTensor],
+    target_snrs: Union[float, BatchTensor],
     psd: PSDTensor,
     sample_rate: float,
-    highpass: Union[float, TensorType["frequency"], None] = None,
+    highpass: Union[float, Float[Tensor, " frequency"], None] = None,
 ) -> WaveformTensor:
     """Scale interferometer responses such that they have a desired SNR
 

ml4gw/nn/autoencoder/base.py

@@ -1,7 +1,8 @@
 from collections.abc import Sequence
-from typing import Optional
+from typing import Optional, Tuple, Union
 
 import torch
+from torch import Tensor
 
 from ml4gw.nn.autoencoder.skip_connection import SkipConnection
 
@@ -27,12 +28,16 @@ class Autoencoder(torch.nn.Module):
     and how they operate.
     """
 
-    def __init__(self, skip_connection: Optional[SkipConnection] = None):
+    def __init__(
+        self, skip_connection: Optional[SkipConnection] = None
+    ) -> None:
         super().__init__()
         self.skip_connection = skip_connection
         self.blocks = torch.nn.ModuleList()
 
-    def encode(self, *X: torch.Tensor, return_states: bool = False):
+    def encode(
+        self, *X: Tensor, return_states: bool = False
+    ) -> Union[Tensor, Tuple[Tensor, Sequence]]:
         states = []
         for block in self.blocks:
             if isinstance(X, tuple):
@@ -48,7 +53,7 @@ class Autoencoder(torch.nn.Module):
             return X, states[:-1]
         return X
 
-    def decode(self, *X, states: Optional[Sequence[torch.Tensor]] = None):
+    def decode(self, *X, states: Optional[Sequence[Tensor]] = None) -> Tensor:
         if self.skip_connection is not None and states is None:
             raise ValueError(
                 "Must pass intermediate states when autoencoder "
@@ -76,7 +81,7 @@ class Autoencoder(torch.nn.Module):
                 X = self.skip_connection(X, state)
         return X
 
-    def forward(self, *X):
+    def forward(self, *X: Tensor) -> Tensor:
         return_states = self.skip_connection is not None
         X = self.encode(*X, return_states=return_states)
         if return_states:
@@ -84,6 +89,6 @@ class Autoencoder(torch.nn.Module):
         else:
             states = None
 
-        if isinstance(X, torch.Tensor):
+        if isinstance(X, Tensor):
             X = (X,)
         return self.decode(*X, states=states)

ml4gw/nn/autoencoder/convolutional.py

@@ -2,6 +2,7 @@ from collections.abc import Callable, Sequence
 from typing import Optional
 
 import torch
+from torch import Tensor
 
 from ml4gw.nn.autoencoder.base import Autoencoder
 from ml4gw.nn.autoencoder.skip_connection import SkipConnection
@@ -64,12 +65,12 @@ class ConvBlock(Autoencoder):
         self.encode_norm = norm(out_channels)
         self.decode_norm = norm(decode_channels)
 
-    def encode(self, X):
+    def encode(self, X: Tensor) -> Tensor:
         X = self.encode_layer(X)
         X = self.encode_norm(X)
         return self.activation(X)
 
-    def decode(self, X):
+    def decode(self, X: Tensor) -> Tensor:
         X = self.decode_layer(X)
         X = self.decode_norm(X)
         return self.output_activation(X)
@@ -144,13 +145,15 @@ class ConvolutionalAutoencoder(Autoencoder):
             self.blocks.append(block)
             in_channels = channels * groups
 
-    def decode(self, *X, states=None, input_size: Optional[int] = None):
+    def decode(
+        self, *X, states=None, input_size: Optional[int] = None
+    ) -> Tensor:
         X = super().decode(*X, states=states)
         if input_size is not None:
             return match_size(X, input_size)
         return X
 
-    def forward(self, X):
+    def forward(self, X: Tensor) -> Tensor:
         input_size = X.size(-1)
         X = super().forward(X)
         return match_size(X, input_size)

ml4gw/nn/autoencoder/skip_connection.py

@@ -1,31 +1,32 @@
 import torch
+from torch import Tensor
 
 from ml4gw.nn.autoencoder.utils import match_size
 
 
 class SkipConnection(torch.nn.Module):
-    def forward(self, X: torch.Tensor, state: torch.Tensor):
+    def forward(self, X: Tensor, state: Tensor) -> Tensor:
         return match_size(X, state.size(-1))
 
-    def get_out_channels(self, in_channels):
+    def get_out_channels(self, in_channels: int) -> int:
         return in_channels
 
 
 class AddSkipConnect(SkipConnection):
-    def forward(self, X, state):
+    def forward(self, X: Tensor, state: Tensor) -> Tensor:
         X = super().forward(X, state)
         return X + state
 
 
 class ConcatSkipConnect(SkipConnection):
-    def __init__(self, groups: int = 1):
+    def __init__(self, groups: int = 1) -> None:
         super().__init__()
         self.groups = groups
 
-    def get_out_channels(self, in_channels):
+    def get_out_channels(self, in_channels: int) -> int:
         return 2 * in_channels
 
-    def forward(self, X, state):
+    def forward(self, X: Tensor, state: Tensor) -> Tensor:
         X = super().forward(X, state)
         if self.groups == 1:
             return torch.cat([X, state], dim=1)

ml4gw/nn/autoencoder/utils.py

@@ -1,7 +1,8 @@
 import torch
+from torch import Tensor
 
 
-def match_size(X: torch.Tensor, target_size: int):
+def match_size(X: Tensor, target_size: int) -> Tensor:
     diff = target_size - X.size(-1)
     left = int(diff // 2)
     right = diff - left

ml4gw/nn/norm.py

@@ -1,6 +1,8 @@
 from typing import Callable, Optional
 
 import torch
+from jaxtyping import Float
+from torch import Tensor
 
 NormLayer = Callable[[int], torch.nn.Module]
 
@@ -31,7 +33,9 @@ class GroupNorm1D(torch.nn.Module):
         self.weight = torch.nn.Parameter(torch.ones(shape))
         self.bias = torch.nn.Parameter(torch.zeros(shape))
 
-    def forward(self, x):
+    def forward(
+        self, x: Float[Tensor, "batch channel length"]
+    ) -> Float[Tensor, "batch channel length"]:
         if len(x.shape) != 3:
             raise ValueError(
                 "GroupNorm1D requires 3-dimensional input, "

ml4gw/nn/streaming/online_average.py

@@ -1,11 +1,11 @@
 from typing import Optional, Tuple
 
 import torch
+from jaxtyping import Float
+from torch import Tensor
 
 from ml4gw.utils.slicing import unfold_windows
 
-Tensor = torch.Tensor
-
 
 class OnlineAverager(torch.nn.Module):
     """
@@ -70,12 +70,14 @@ class OnlineAverager(torch.nn.Module):
         weights = unfold_windows(weights, weight_size, update_size)
         self.register_buffer("weights", weights)
 
-    def get_initial_state(self):
+    def get_initial_state(self) -> Float[Tensor, "channel time"]:
         return torch.zeros((self.num_channels, self.state_size))
 
     def forward(
-        self, update: torch.Tensor, state: Optional[torch.Tensor] = None
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        self,
+        update: Float[Tensor, "batch channel time1"],
+        state: Optional[Float[Tensor, "channel time2"]] = None,
+    ) -> Tuple[Float[Tensor, "channel time3"], Float[Tensor, "channel time4"]]:
         if state is None:
             state = self.get_initial_state()
 

ml4gw/nn/streaming/snapshotter.py

@@ -1,6 +1,8 @@
 from typing import Optional, Sequence, Tuple
 
 import torch
+from jaxtyping import Float
+from torch import Tensor
 
 from ml4gw.utils.slicing import unfold_windows
 
@@ -82,14 +84,14 @@ class Snapshotter(torch.nn.Module):
         self.channels_per_snapshot = channels_per_snapshot
         self.num_channels = num_channels
 
-    def get_initial_state(self):
+    def get_initial_state(self) -> Float[Tensor, "channel time"]:
         return torch.zeros((self.num_channels, self.state_size))
 
-    # TODO: use torchtyping annotations to make
-    # clear what the expected shapes are
     def forward(
-        self, update: torch.Tensor, snapshot: Optional[torch.Tensor] = None
-    ) -> Tuple[torch.Tensor, ...]:
+        self,
+        update: Float[Tensor, "channel time1"],
+        snapshot: Optional[Float[Tensor, "channel time2"]] = None,
+    ) -> Tuple[Tensor, ...]:
         if snapshot is None:
             snapshot = self.get_initial_state()
 

ml4gw/spectral.py

@@ -12,14 +12,18 @@ https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.csd.html
 from typing import Optional, Union
 
 import torch
-from torchtyping import TensorType
+from jaxtyping import Float
+from torch import Tensor
 
-from ml4gw import types
+from ml4gw.types import (
+    FrequencySeries1to3d,
+    PSDTensor,
+    TimeSeries1to3d,
+    WaveformTensor,
+)
 
-time = None
 
-
-def median(x, axis):
+def median(x: Float[Tensor, "... size"], axis: int) -> Float[Tensor, "..."]:
     """
     Implements a median calculation that matches numpy's
     behavior for an even number of elements and includes
@@ -33,7 +37,7 @@ def median(x, axis):
 
 
 def _validate_shapes(
-    x: torch.Tensor, nperseg: int, y: Optional[torch.Tensor] = None
+    x: Tensor, nperseg: int, y: Optional[Tensor] = None
 ) -> None:
     if x.shape[-1] < nperseg:
         raise ValueError(
@@ -83,14 +87,14 @@ def _validate_shapes(
 
 
 def fast_spectral_density(
-    x: torch.Tensor,
+    x: TimeSeries1to3d,
     nperseg: int,
     nstride: int,
-    window: torch.Tensor,
-    scale: torch.Tensor,
+    window: Float[Tensor, " {nperseg//2+1}"],
+    scale: float,
     average: str = "median",
-    y: Optional[torch.Tensor] = None,
-) -> torch.Tensor:
+    y: Optional[TimeSeries1to3d] = None,
+) -> FrequencySeries1to3d:
     """
     Compute the power spectral density of a multichannel
     timeseries or a batch of multichannel timeseries, or
@@ -107,9 +111,9 @@ def fast_spectral_density(
             The timeseries tensor whose power spectral density
             to compute, or for cross spectral density the
             timeseries whose fft will be conjugated. Can have
-            shape either
-            `(batch_size, num_channels, length * sample_rate)`
-            or `(num_channels, length * sample_rate)`.
+            shape `(batch_size, num_channels, length * sample_rate)`,
+            `(num_channels, length * sample_rate)`, or
+            `(length * sample_rate)`.
         nperseg:
             Number of samples included in each FFT window
         nstride:
@@ -150,7 +154,7 @@ def fast_spectral_density(
             channel in `x` across _all_ of `x`'s batch elements.
     Returns:
         Tensor of power spectral densities of `x` or its cross spectral
-            density with the timeseries in `y`.
+        density with the timeseries in `y`.
     """
 
     _validate_shapes(x, nperseg, y)
@@ -240,17 +244,16 @@ def fast_spectral_density(
 
 
 def spectral_density(
-    x: torch.Tensor,
+    x: TimeSeries1to3d,
     nperseg: int,
     nstride: int,
-    window: torch.Tensor,
-    scale: torch.Tensor,
+    window: Float[Tensor, " {nperseg//2+1}"],
+    scale: float,
     average: str = "median",
-) -> torch.Tensor:
+) -> FrequencySeries1to3d:
     """
     Compute the power spectral density of a multichannel
-    timeseries or a batch of multichannel timeseries, or
-    the cross power spectral density of two such timeseries.
+    timeseries or a batch of multichannel timeseries.
     This implementation is exact for all frequency bins, but
     slower than the fast implementation.
 
@@ -259,9 +262,9 @@ def spectral_density(
             The timeseries tensor whose power spectral density
             to compute, or for cross spectral density the
             timeseries whose fft will be conjugated. Can have
-            shape either
-            `(batch_size, num_channels, length * sample_rate)`
-            or `(num_channels, length * sample_rate)`.
+            shape `(batch_size, num_channels, length * sample_rate)`,
+            `(num_channels, length * sample_rate)`, or
+            `(length * sample_rate)`.
         nperseg:
             Number of samples included in each FFT window
         nstride:
@@ -336,11 +339,11 @@ def spectral_density(
 
 
 def truncate_inverse_power_spectrum(
-    psd: types.PSDTensor,
-    fduration: Union[TensorType["time"], float],
+    psd: PSDTensor,
+    fduration: Union[Float[Tensor, " time"], float],
     sample_rate: float,
     highpass: Optional[float] = None,
-) -> types.PSDTensor:
+) -> PSDTensor:
     """
     Truncate the length of the time domain response
     of a whitening filter built using the specified
@@ -399,7 +402,7 @@ def truncate_inverse_power_spectrum(
     q = torch.fft.irfft(inv_asd, n=N, norm="forward", dim=-1)
 
     # taper the edges of the TD filter
-    if isinstance(fduration, torch.Tensor):
+    if isinstance(fduration, Tensor):
         pad = fduration.size(-1) // 2
         window = fduration
     else:
@@ -422,8 +425,8 @@ def truncate_inverse_power_spectrum(
 
 
 def normalize_by_psd(
-    X: types.WaveformTensor,
-    psd: types.PSDTensor,
+    X: WaveformTensor,
+    psd: PSDTensor,
     sample_rate: float,
     pad: int,
 ):
@@ -447,12 +450,12 @@ def normalize_by_psd(
 
 
 def whiten(
-    X: types.WaveformTensor,
-    psd: types.PSDTensor,
-    fduration: Union[TensorType["time"], float],
+    X: WaveformTensor,
+    psd: PSDTensor,
+    fduration: Union[Float[Tensor, " time"], float],
     sample_rate: float,
     highpass: Optional[float] = None,
-) -> types.WaveformTensor:
+) -> WaveformTensor:
     """
     Whiten a batch of timeseries using the specified
     background one-sided power spectral densities (PSDs),
@@ -460,7 +463,8 @@ def whiten(
     `fduration` and possibly to highpass filter.
 
     Args:
-        X: batch of multichannel timeseries to whiten
+        X:
+            batch of multichannel timeseries to whiten
         psd:
             PSDs use to whiten the data. The frequency
             response of the whitening filter will be roughly
@@ -496,7 +500,7 @@ def whiten(
 
     # figure out how much data we'll need to slice
     # off after whitening
-    if isinstance(fduration, torch.Tensor):
+    if isinstance(fduration, Tensor):
         pad = fduration.size(-1) // 2
     else:
         pad = int(fduration * sample_rate / 2)