ml4gw 0.4.2__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/constants.py

@@ -0,0 +1,45 @@
+"""
+Various constants, all in SI units.
+"""
+
+EulerGamma = 0.577215664901532860606512090082402431
+
+MSUN = 1.988409902147041637325262574352366540e30  # kg
+"""Solar mass"""
+
+MRSUN = 1.476625038050124729627979840144936351e3
+"""Geometrized nominal solar mass, m"""
+
+G = 6.67430e-11  # m^3 / kg / s^2
+"""Newton's gravitational constant"""
+
+C = 299792458.0  # m / s
+"""Speed of light"""
+
+"""Pi"""
+PI = 3.141592653589793238462643383279502884
+
+TWO_PI = 6.283185307179586476925286766559005768
+
+gt = G * MSUN / (C**3.0)
+"""
+G MSUN / C^3 in seconds
+"""
+
+MTSUN_SI = 4.925490947641266978197229498498379006e-6
+"""1 solar mass in seconds. Same value as lal.MTSUN_SI"""
+
+m_per_Mpc = 3.085677581491367278913937957796471611e22
+"""
+Meters per Mpc.
+"""
+
+MPC_SEC = m_per_Mpc / C
+"""
+1 Mpc in seconds.
+"""
+
+clightGpc = C / 3.0856778570831e22
+"""
+Speed of light in vacuum (:math:`c`), in gigaparsecs per second
+"""

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

@@ -4,11 +4,13 @@ from specified distributions. Each callable should map from
 an integer `N` to a 1D torch `Tensor` containing `N` samples
 from the corresponding distribution.
 """
-
+import math
 from typing import Optional
 
 import torch
 import torch.distributions as dist
+from jaxtyping import Float
+from torch import Tensor
 
 
 class Cosine(dist.Distribution):
@@ -21,20 +23,21 @@ class Cosine(dist.Distribution):
 
     def __init__(
         self,
-        low: float = torch.as_tensor(-torch.pi / 2),
-        high: float = torch.as_tensor(torch.pi / 2),
+        low: float = -math.pi / 2,
+        high: float = math.pi / 2,
         validate_args=None,
     ):
         batch_shape = torch.Size()
         super().__init__(batch_shape, validate_args=validate_args)
-        self.low = low
-        self.norm = 1 / (torch.sin(high) - torch.sin(low))
+        self.low = torch.as_tensor(low)
+        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
@@ -48,13 +51,16 @@ class Sine(dist.TransformedDistribution):
 
     def __init__(
         self,
-        low: float = torch.as_tensor(0),
-        high: float = torch.as_tensor(torch.pi),
+        low: float = 0.0,
+        high: float = math.pi,
         validate_args=None,
     ):
+        low = torch.as_tensor(low)
+        high = torch.as_tensor(high)
         base_dist = Cosine(
             low - torch.pi / 2, high - torch.pi / 2, validate_args
         )
+
         super().__init__(
             base_dist,
             [
@@ -153,14 +159,14 @@ class DeltaFunction(dist.Distribution):
 
     def __init__(
         self,
-        peak: float = torch.as_tensor(0.0),
+        peak: float = 0.0,
         validate_args=None,
     ):
         batch_shape = torch.Size()
         super().__init__(batch_shape, validate_args=validate_args)
-        self.peak = peak
+        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,15 @@ 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, "
+                f"received {len(x.shape)} dimensional input"
+            )
+
         keepdims = self.num_groups == self.num_channels
 
         # compute group variance via the E[x**2] - E**2[x] trick

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()