v0.5.0 (#146)

* re-implement distributions using torch.distributions (#110) * loosely re-implement distributions using torch.distributions * further enhancement * rename * log_uniform as transformed dist * delete uniform and log_normal since they are already implemented * bring back LogNormal, change parameter sampler * add delta function distribution, minor fixes to other distributions * fix * Initial commit of q-transform * Added a MultiQTransform * Reparameterized to number of t and f bins * Re-factored SingleQTransform to allow for eventually using different q's for batch/channels * Changed interpolation to better match gwpy and added option to specify frequency range in which to search for max energy tile * Added QScan to __init__ * Changed from torch median to quantile to match numpy median * Added documentation to qtransform * Added gwpy to dev dependencies * Added tests and corrected get_freqs bug * Changed method of normalization and updated documentation * Updated more documentation and changed interpolation method * Changed how interpolation shape is parameterized * Reverted documentation to state expectation of 3D input * Fixed type hint for spectrogram_shape in qtransform * Switched tuple to Tuple * `InMemoryDataset` improvements (#119) * in memory dataset inherits from torch iterable dataset * pre-commit issues * fix yielding logic * fix in memory dataset tests * handle deprecated transpose ops in phenomd (#120) * Torch dependency fix into dev (#124) * udpate poetry lock * increment version once more * re add gwpy dev dep * update poetry lock * fix poetry conflic * fix poety conflict * Poetry fix (#126) * remove gwpy dep for now * poetry lock file * Merge `main` into `dev` (#128) * Adding funding statement * Add requirements to docs (#107) * add requirements to docs * remove myst_parser req * Update requirements.txt * Update conf.py * Update docs.yaml (#108) * Update docs.yaml * Update docs.yaml (#109) * simplify torch dependencies (#122) * simplify torch dependency * remove python constratints * fix pre-commit * Merge `dev` into `main` for 0.4.2 release (#123) * re-implement distributions using torch.distributions (#110) * loosely re-implement distributions using torch.distributions * further enhancement * rename * log_uniform as transformed dist * delete uniform and log_normal since they are already implemented * bring back LogNormal, change parameter sampler * add delta function distribution, minor fixes to other distributions * fix * Initial commit of q-transform * Added a MultiQTransform * Reparameterized to number of t and f bins * Re-factored SingleQTransform to allow for eventually using different q's for batch/channels * Changed interpolation to better match gwpy and added option to specify frequency range in which to search for max energy tile * Added QScan to __init__ * Changed from torch median to quantile to match numpy median * Added documentation to qtransform * Added gwpy to dev dependencies * Added tests and corrected get_freqs bug * Changed method of normalization and updated documentation * Updated more documentation and changed interpolation method * Changed how interpolation shape is parameterized * Reverted documentation to state expectation of 3D input * Fixed type hint for spectrogram_shape in qtransform * Switched tuple to Tuple * `InMemoryDataset` improvements (#119) * in memory dataset inherits from torch iterable dataset * pre-commit issues * fix yielding logic * fix in memory dataset tests * handle deprecated transpose ops in phenomd (#120) * Torch dependency fix into dev (#124) * udpate poetry lock * increment version once more * re add gwpy dev dep * update poetry lock * fix poetry conflic * fix poety conflict * Poetry fix (#126) * remove gwpy dep for now * poetry lock file --------- Co-authored-by: Deep Chatterjee <[email protected]> Co-authored-by: wbenoit26 <[email protected]> Co-authored-by: William Benoit <[email protected]> Co-authored-by: William Benoit <[email protected]> * increment version (#127) --------- Co-authored-by: wbenoit26 <[email protected]> Co-authored-by: Deep Chatterjee <[email protected]> Co-authored-by: William Benoit <[email protected]> Co-authored-by: William Benoit <[email protected]> * Fix Distributions module typehints (#129) * fix type hints in dists * fix type hints in dists * Added ringdown waveform function * fix prefactor four thirds * Ringdown Waveform updated * Updated Ringdown with added precision and equation changes * Added KAGRA geometry * Add test for K1 geometry * Fix pre-commit issue * Restrict numpy * Unrestrict numpy * Update tox settings * Added GroupNorm1D dimension check * fix thinko in variable names (#135) * Refactor TaylorF2 and IMRPhenomD (#137) * refactor taylorf2 and phenom_d into a torch module * subtract 3PN terms from phasing and Dphasing * update tolerance for phenomd tests * Added ringdown * fixed negative time issues * Add IMRPhenomPv2 waveform (#138) * add IMRPhenomPv2 torch implementation * perform phase correction term by taking finite difference derivative * add tests for phenomPv2 * added constants * use dtype of parameters for sg waveform * Fixed constants import * Allow arbitrary window in `SpectralDensity` (#140) * allow arbitrary window function * allow arbitrary window function * update error message * revert spectral transform filename * delete test_spectral * constrain numpy version to <2.0.0 (#145) * Use .mT instead of .T (#147) * update to use .mT instead of .T * update taylorf2 to use .mT * Docstring for phenom_d, phenom_p, tf2 (#148) * docstring for phenom_d, phenom_p, tf2 * update docs for consistency * add units to some params * forward/__call__ and cross/plus refactor (#149) * remove .forward and call directly * update sine gaussian to use forward instead of __call__ * update ringdown to use forward instead of __call__ * swap cross and plus polarisations in return. now returns (cross, plus) * Resolve main conflicts (#151) * Adding funding statement * Add requirements to docs (#107) * add requirements to docs * remove myst_parser req * Update requirements.txt * Update conf.py * Update docs.yaml (#108) * Update docs.yaml * Update docs.yaml (#109) * simplify torch dependencies (#122) * simplify torch dependency * remove python constratints * fix pre-commit * Merge `dev` into `main` for 0.4.2 release (#123) * re-implement distributions using torch.distributions (#110) * loosely re-implement distributions using torch.distributions * further enhancement * rename * log_uniform as transformed dist * delete uniform and log_normal since they are already implemented * bring back LogNormal, change parameter sampler * add delta function distribution, minor fixes to other distributions * fix * Initial commit of q-transform * Added a MultiQTransform * Reparameterized to number of t and f bins * Re-factored SingleQTransform to allow for eventually using different q's for batch/channels * Changed interpolation to better match gwpy and added option to specify frequency range in which to search for max energy tile * Added QScan to __init__ * Changed from torch median to quantile to match numpy median * Added documentation to qtransform * Added gwpy to dev dependencies * Added tests and corrected get_freqs bug * Changed method of normalization and updated documentation * Updated more documentation and changed interpolation method * Changed how interpolation shape is parameterized * Reverted documentation to state expectation of 3D input * Fixed type hint for spectrogram_shape in qtransform * Switched tuple to Tuple * `InMemoryDataset` improvements (#119) * in memory dataset inherits from torch iterable dataset * pre-commit issues * fix yielding logic * fix in memory dataset tests * handle deprecated transpose ops in phenomd (#120) * Torch dependency fix into dev (#124) * udpate poetry lock * increment version once more * re add gwpy dev dep * update poetry lock * fix poetry conflic * fix poety conflict * Poetry fix (#126) * remove gwpy dep for now * poetry lock file --------- Co-authored-by: Deep Chatterjee <[email protected]> Co-authored-by: wbenoit26 <[email protected]> Co-authored-by: William Benoit <[email protected]> Co-authored-by: William Benoit <[email protected]> * increment version (#127) --------- Co-authored-by: wbenoit26 <[email protected]> Co-authored-by: Deep Chatterjee <[email protected]> Co-authored-by: William Benoit <[email protected]> Co-authored-by: William Benoit <[email protected]> * resolve merge conflicts * resolve merge conflicts --------- Co-authored-by: Deep Chatterjee <[email protected]> Co-authored-by: wbenoit26 <[email protected]> Co-authored-by: William Benoit <[email protected]> Co-authored-by: William Benoit <[email protected]> Co-authored-by: Trevor William Blodgett <[email protected]> Co-authored-by: Ravi Kumar <[email protected]> Co-authored-by: Trevor William Blodgett <[email protected]> Co-authored-by: Ravi Kumar <[email protected]>
ML4GW · Jul 15, 2024 · c2e1937 · c2e1937
1 parent 74c13d5
commit c2e1937
Show file tree

Hide file tree

Showing 18 changed files with 3,302 additions and 2,072 deletions.
diff --git a/ml4gw/constants.py b/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
+"""
diff --git a/ml4gw/distributions.py b/ml4gw/distributions.py
@@ -4,7 +4,7 @@
 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(

diff --git a/ml4gw/nn/norm.py b/ml4gw/nn/norm.py
@@ -32,6 +32,12 @@ def __init__(
         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

diff --git a/ml4gw/transforms/spectral.py b/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 @@ def __init__(
         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 @@ def __init__(
         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

diff --git a/ml4gw/utils/interferometer.py b/ml4gw/utils/interferometer.py
@@ -35,6 +35,12 @@ def __init__(self, name: str):
             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, "

diff --git a/ml4gw/waveforms/__init__.py b/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