ml4gw 0.4.1__tar.gz → 0.5.0__tar.gz

{ml4gw-0.4.1 → ml4gw-0.5.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ml4gw
-Version: 0.4.1
+Version: 0.5.0
 Summary: Tools for training torch models on gravitational wave data
 Author: Alec Gunny
 Author-email: alec.gunny@ligo.org
@@ -10,6 +10,7 @@ Classifier: Programming Language :: Python :: 3.8
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
+Requires-Dist: numpy (<2.0.0)
 Requires-Dist: torch (>=2.0,<3.0)
 Requires-Dist: torchaudio (>=2.0,<3.0)
 Requires-Dist: torchtyping (>=0.1,<0.2)

ml4gw-0.5.0/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-0.4.1 → ml4gw-0.5.0}/ml4gw/dataloading/in_memory_dataset.py

@@ -7,7 +7,7 @@ from ml4gw import types
 from ml4gw.utils.slicing import slice_kernels
 
 
-class InMemoryDataset:
+class InMemoryDataset(torch.utils.data.IterableDataset):
     """Dataset for iterating through in-memory multi-channel timeseries
 
     Dataset for arrays of timeseries data which can be stored
@@ -131,7 +131,6 @@ class InMemoryDataset:
         self.batches_per_epoch = batches_per_epoch
         self.shuffle = shuffle
         self.coincident = coincident
-        self._i = self._idx = None
 
     @property
     def num_kernels(self) -> int:
@@ -157,7 +156,7 @@ class InMemoryDataset:
             num_kernels = self.num_kernels ** len(self.X)
             return (num_kernels - 1) // self.batch_size + 1
 
-    def __iter__(self):
+    def init_indices(self):
         """
         Initialize arrays of indices we'll use to slice
         through X and y at iteration time. This helps by
@@ -204,36 +203,23 @@ class InMemoryDataset:
             # the simplest case: deteriminstic and coincident
             idx = torch.arange(num_kernels, device=device)
 
-        self._idx = idx
-        self._i = 0
-        return self
+        return idx
 
-    def __next__(
+    def __iter__(
         self,
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
-        if self._i is None or self._idx is None:
-            raise TypeError(
-                "Must initialize InMemoryDataset iteration "
-                "before calling __next__"
-            )
 
-        # check if we're out of batches, and if so
-        # make sure to reset before stopping
-        if self._i >= len(self):
-            self._i = self._idx = None
-            raise StopIteration
-
-        # slice the array of _indices_ we'll be using to
-        # slice our timeseries, and scale them by the stride
-        slc = slice(self._i * self.batch_size, (self._i + 1) * self.batch_size)
-        idx = self._idx[slc] * self.stride
-
-        # slice our timeseries
-        X = slice_kernels(self.X, idx, self.kernel_size)
-        if self.y is not None:
-            y = slice_kernels(self.y, idx, self.kernel_size)
-
-        self._i += 1
-        if self.y is not None:
-            return X, y
-        return X
+        indices = self.init_indices()
+        for i in range(len(self)):
+            # slice the array of _indices_ we'll be using to
+            # slice our timeseries, and scale them by the stride
+            slc = slice(i * self.batch_size, (i + 1) * self.batch_size)
+            idx = indices[slc] * self.stride
+
+            # slice our timeseries
+            X = slice_kernels(self.X, idx, self.kernel_size)
+            if self.y is not None:
+                y = slice_kernels(self.y, idx, self.kernel_size)
+                yield X, y
+            else:
+                yield X

ml4gw-0.5.0/ml4gw/distributions.py

@@ -0,0 +1,170 @@
+"""
+Module containing callables classes for generating samples
+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
+
+
+class Cosine(dist.Distribution):
+    """
+    Cosine distribution based on
+    ``torch.distributions.TransformedDistribution``.
+    """
+
+    arg_constraints = {}
+
+    def __init__(
+        self,
+        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 = 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:
+        u = torch.rand(sample_shape, device=self.low.device)
+        return torch.arcsin(u / self.norm + torch.sin(self.low))
+
+    def log_prob(self, value):
+        value = torch.as_tensor(value)
+        inside_range = (value >= self.low) & (value <= self.high)
+        return value.cos().log() * inside_range
+
+
+class Sine(dist.TransformedDistribution):
+    """
+    Sine distribution based on
+    ``torch.distributions.TransformedDistribution``.
+    """
+
+    def __init__(
+        self,
+        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,
+            [
+                dist.AffineTransform(
+                    loc=torch.pi / 2,
+                    scale=1,
+                )
+            ],
+            validate_args=validate_args,
+        )
+
+
+class LogUniform(dist.TransformedDistribution):
+    """
+    Sample from a log uniform distribution
+    """
+
+    def __init__(self, low: float, high: float, validate_args=None):
+        base_dist = dist.Uniform(
+            torch.as_tensor(low).log(),
+            torch.as_tensor(high).log(),
+            validate_args,
+        )
+        super().__init__(
+            base_dist,
+            [dist.ExpTransform()],
+            validate_args=validate_args,
+        )
+
+
+class LogNormal(dist.LogNormal):
+    def __init__(
+        self,
+        mean: float,
+        std: float,
+        low: Optional[float] = None,
+        validate_args=None,
+    ):
+        self.low = low
+        super().__init__(loc=mean, scale=std, validate_args=validate_args)
+
+    def support(self):
+        if self.low is not None:
+            return dist.constraints.greater_than(self.low)
+
+
+class PowerLaw(dist.TransformedDistribution):
+    """
+    Sample from a power law distribution,
+    .. math::
+        p(x) \approx x^{\alpha}.
+
+    Index alpha cannot be 0, since it is equivalent to a Uniform distribution.
+    This could be used, for example, as a universal distribution of
+    signal-to-noise ratios (SNRs) from uniformly volume distributed
+    sources
+    .. math::
+
+       p(\rho) = 3*\rho_0^3 / \rho^4
+
+    where :math:`\rho_0` is a representative minimum SNR
+    considered for detection. See, for example,
+    `Schutz (2011) <https://arxiv.org/abs/1102.5421>`_.
+    Or, for example, ``index=2`` for uniform in Euclidean volume.
+    """
+
+    support = dist.constraints.nonnegative
+
+    def __init__(
+        self, minimum: float, maximum: float, index: int, validate_args=None
+    ):
+        if index == 0:
+            raise RuntimeError("Index of 0 is the same as Uniform")
+        elif index == -1:
+            base_min = torch.as_tensor(minimum).log()
+            base_max = torch.as_tensor(maximum).log()
+            transforms = [dist.ExpTransform()]
+        else:
+            index_plus = index + 1
+            base_min = minimum**index_plus / index_plus
+            base_max = maximum**index_plus / index_plus
+            transforms = [
+                dist.AffineTransform(loc=0, scale=index_plus),
+                dist.PowerTransform(1 / index_plus),
+            ]
+        base_dist = dist.Uniform(base_min, base_max, validate_args=False)
+        super().__init__(
+            base_dist,
+            transforms,
+            validate_args=validate_args,
+        )
+
+
+class DeltaFunction(dist.Distribution):
+    arg_constraints = {}
+
+    def __init__(
+        self,
+        peak: float = 0.0,
+        validate_args=None,
+    ):
+        batch_shape = torch.Size()
+        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:
+        return self.peak * torch.ones(
+            sample_shape, device=self.peak.device, dtype=torch.float32
+        )

{ml4gw-0.4.1 → ml4gw-0.5.0}/ml4gw/nn/norm.py

@@ -32,6 +32,12 @@ class GroupNorm1D(torch.nn.Module):
         self.bias = torch.nn.Parameter(torch.zeros(shape))
 
     def forward(self, x):
+        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-0.4.1 → ml4gw-0.5.0}/ml4gw/transforms/__init__.py

@@ -1,4 +1,5 @@
 from .pearson import ShiftedPearsonCorrelation
+from .qtransform import QScan, SingleQTransform
 from .scaler import ChannelWiseScaler
 from .snr_rescaler import SnrRescaler
 from .spectral import SpectralDensity