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

@@ -4,7 +4,7 @@ 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
@@ -21,14 +21,15 @@ 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:
         u = torch.rand(sample_shape, device=self.low.device)
@@ -48,13 +49,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,12 +157,12 @@ 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:
         return self.peak * torch.ones(

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

@@ -34,6 +34,10 @@ class SpectralDensity(torch.nn.Module):
         average:
             Aggregation method to use for combining windowed FFTs.
             Allowed values are `"mean"` and `"median"`.
+        window:
+            Window array to multiply by each FFT window before
+            FFT computation. Should have length `nperseg`.
+            Defaults to a hanning window.
         fast:
             Whether to use a faster spectral density computation that
             support cross spectral density, or a slower one which does
@@ -47,6 +51,7 @@ class SpectralDensity(torch.nn.Module):
         fftlength: float,
         overlap: Optional[float] = None,
         average: str = "mean",
+        window: Optional[torch.Tensor] = None,
         fast: bool = False,
     ) -> None:
         if overlap is None:
@@ -63,11 +68,18 @@ class SpectralDensity(torch.nn.Module):
         self.nperseg = int(fftlength * sample_rate)
         self.nstride = self.nperseg - int(overlap * sample_rate)
 
-        # TODOs: Do we allow for arbitrary windows?
-        # Making this buffer persistent in case we want
-        # to implement this down the line, so that custom
-        # windows can be loaded in.
-        self.register_buffer("window", torch.hann_window(self.nperseg))
+        # if no window is provided, default to a hanning window;
+        # validate that window is correct size
+        if window is None:
+            window = torch.hann_window(self.nperseg)
+
+        if window.size(0) != self.nperseg:
+            raise ValueError(
+                "Window must have length {} got {}".format(
+                    self.nperseg, window.size(0)
+                )
+            )
+        self.register_buffer("window", window)
 
         # scale corresponds to "density" normalization, worth
         # considering adding this as a kwarg and changing this calc

ml4gw/utils/interferometer.py

@@ -35,6 +35,12 @@ class InterferometerGeometry:
             self.vertex = torch.Tensor(
                 (4.54637409900e06, 8.42989697626e05, 4.37857696241e06)
             )
+        elif name == "K1":
+            self.x_arm = torch.Tensor((-0.3759040, -0.8361583, 0.3994189))
+            self.y_arm = torch.Tensor((0.7164378, 0.01114076, 0.6975620))
+            self.vertex = torch.Tensor(
+                (-3777336.024, 3484898.411, 3765313.697)
+            )
         else:
             raise ValueError(
                 f"{name} is not recognized as an interferometer, "

ml4gw/waveforms/__init__.py

@@ -1,3 +1,5 @@
 from .phenom_d import IMRPhenomD
+from .phenom_p import IMRPhenomPv2
+from .ringdown import Ringdown
 from .sine_gaussian import SineGaussian
 from .taylorf2 import TaylorF2