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

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/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/distributions.py

@@ -4,94 +4,114 @@ 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 Uniform:
+class Cosine(dist.Distribution):
     """
-    Sample uniformly between `low` and `high`.
+    Cosine distribution based on
+    ``torch.distributions.TransformedDistribution``.
     """
 
-    def __init__(self, low: float = 0, high: float = 1) -> None:
-        self.low = low
-        self.high = high
-
-    def __call__(self, N: int) -> torch.Tensor:
-        return self.low + torch.rand(size=(N,)) * (self.high - self.low)
+    arg_constraints = {}
 
-
-class Cosine:
+    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):
     """
-    Sample from a raised Cosine distribution between
-    `low` and `high`. Based on the implementation from
-    bilby documented here:
-    https://lscsoft.docs.ligo.org/bilby/api/bilby.core.prior.analytical.Cosine.html  # noqa
+    Sine distribution based on
+    ``torch.distributions.TransformedDistribution``.
     """
 
     def __init__(
-        self, low: float = -math.pi / 2, high: float = math.pi / 2
-    ) -> None:
-        self.low = low
-        self.norm = 1 / (math.sin(high) - math.sin(low))
-
-    def __call__(self, N: int) -> torch.Tensor:
-        """
-        Implementation lifted from
-        https://lscsoft.docs.ligo.org/bilby/_modules/bilby/core/prior/analytical.html#Cosine # noqa
-        """
-        u = torch.rand(size=(N,))
-        return torch.arcsin(u / self.norm + math.sin(self.low))
-
-
-class LogNormal:
+        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 normal distribution with the
-    specified `mean` and standard deviation `std`.
-    If a `low` value is specified, values sampled
-    lower than this will be clipped to `low`.
+    Sample from a log uniform distribution
     """
 
-    def __init__(
-        self, mean: float, std: float, low: Optional[float] = None
-    ) -> None:
-        self.sigma = math.log((std / mean) ** 2 + 1) ** 0.5
-        self.mu = 2 * math.log(mean / (mean**2 + std**2) ** 0.25)
-        self.low = low
+    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,
+        )
 
-    def __call__(self, N: int) -> torch.Tensor:
 
-        u = self.mu + torch.randn(N) * self.sigma
-        x = torch.exp(u)
+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:
-            x = torch.clip(x, self.low)
-        return x
-
-
-class LogUniform(Uniform):
-    """
-    Sample from a log uniform distribution
-    """
-
-    def __init__(self, low: float, high: float) -> None:
-        super().__init__(math.log(low), math.log(high))
-
-    def __call__(self, N: int) -> torch.Tensor:
-        u = super().__call__(N)
-        return torch.exp(u)
+            return dist.constraints.greater_than(self.low)
 
 
-class PowerLaw:
+class PowerLaw(dist.TransformedDistribution):
     """
     Sample from a power law distribution,
     .. math::
-        p(x) \approx x^{-\alpha}.
+        p(x) \approx x^{\alpha}.
 
-    Index alpha must be greater than 1.
+    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
@@ -102,21 +122,49 @@ class PowerLaw:
     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, x_min: float, x_max: float = float("inf"), alpha: float = 2
-    ) -> None:
-        self.x_min = x_min
-        self.x_max = x_max
-        self.alpha = alpha
-
-        self.normalization = x_min ** (-self.alpha + 1)
-        self.normalization -= x_max ** (-self.alpha + 1)
-
-    def __call__(self, N: int) -> torch.Tensor:
-        u = torch.rand(N)
-        u *= self.normalization
-        samples = self.x_min ** (-self.alpha + 1) - u
-        samples = torch.pow(samples, -1.0 / (self.alpha - 1))
-        return samples
+        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/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/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