ml4gw 0.5.1__tar.gz → 0.6.0__tar.gz

{ml4gw-0.5.1 → ml4gw-0.6.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ml4gw
-Version: 0.5.1
+Version: 0.6.0
 Summary: Tools for training torch models on gravitational wave data
 Author: Alec Gunny
 Author-email: alec.gunny@ligo.org

{ml4gw-0.5.1 → ml4gw-0.6.0}/ml4gw/constants.py

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

{ml4gw-0.5.1 → ml4gw-0.6.0}/ml4gw/spectral.py

@@ -441,7 +441,7 @@ def normalize_by_psd(
 
     # convert back to the time domain and normalize
     # TODO: what's this normalization factor?
-    X = torch.fft.irfft(X_tilde, norm="forward", dim=-1)
+    X = torch.fft.irfft(X_tilde, n=X.shape[-1], norm="forward", dim=-1)
     X = X.float() / sample_rate**0.5
 
     # slice off corrupted data at edges of kernel

{ml4gw-0.5.1 → ml4gw-0.6.0}/ml4gw/transforms/__init__.py

@@ -4,5 +4,6 @@ from .scaler import ChannelWiseScaler
 from .snr_rescaler import SnrRescaler
 from .spectral import SpectralDensity
 from .spectrogram import MultiResolutionSpectrogram
+from .spline_interpolation import SplineInterpolate
 from .waveforms import WaveformProjector, WaveformSampler
 from .whitening import FixedWhiten, Whiten

{ml4gw-0.5.1 → ml4gw-0.6.0}/ml4gw/transforms/qtransform.py

@@ -1,11 +1,13 @@
 import math
-from typing import List, Optional, Tuple
+import warnings
+from typing import List, Tuple
 
 import torch
 import torch.nn.functional as F
 from jaxtyping import Float, Int
 from torch import Tensor
 
+from ml4gw.transforms.spline_interpolation import SplineInterpolate
 from ml4gw.types import FrequencySeries1to3d, TimeSeries1to3d, TimeSeries3d
 
 """
@@ -38,7 +40,6 @@ class QTile(torch.nn.Module):
         mismatch:
             The maximum fractional mismatch between neighboring tiles
 
-
     """
 
     def __init__(
@@ -100,7 +101,9 @@ class QTile(torch.nn.Module):
         ).type(torch.long)
 
     def forward(
-        self, fseries: FrequencySeries1to3d, norm: str = "median"
+        self,
+        fseries: FrequencySeries1to3d,
+        norm: str = "median",
     ) -> TimeSeries1to3d:
         """
         Compute the transform for this row
@@ -144,7 +147,7 @@ class QTile(torch.nn.Module):
                 energy /= means
             else:
                 raise ValueError("Invalid normalisation %r" % norm)
-            return energy.type(torch.float32)
+            energy = energy.type(torch.float32)
         return energy
 
 
@@ -172,6 +175,19 @@ class SingleQTransform(torch.nn.Module):
             be chosen based on q, sample_rate, and duration
         mismatch:
             The maximum fractional mismatch between neighboring tiles
+        interpolation_method:
+            The method by which to interpolate each `QTile` to the specified
+            number of time and frequency bins. The acceptable values are
+            "bilinear", "bicubic", and "spline". The "bilinear" and "bicubic"
+            options will use PyTorch's built-in interpolation modes, while
+            "spline" will use the custom Torch-based implementation in
+            `ml4gw`, as PyTorch does not have spline-based intertpolation.
+            The "spline" mode is most similar to the results of GWpy's
+            Q-transform, which uses `scipy` to do spline interpolation.
+            However, it is also the slowest and most memory intensive due to
+            the matrix equation solving steps. Therefore, the default method
+            is "bicubic" as it produces the most similar results while
+            optimizing for computing performance.
     """
 
     def __init__(
@@ -182,6 +198,7 @@ class SingleQTransform(torch.nn.Module):
         q: float = 12,
         frange: List[float] = [0, torch.inf],
         mismatch: float = 0.2,
+        interpolation_method: str = "bicubic",
     ) -> None:
         super().__init__()
         self.q = q
@@ -190,20 +207,87 @@ class SingleQTransform(torch.nn.Module):
         self.duration = duration
         self.mismatch = mismatch
 
+        # If q is too large, the minimum of the frange computed
+        # below will be larger than the maximum
+        max_q = torch.pi * duration * sample_rate / 50 - 11 ** (0.5)
+        if q >= max_q:
+            raise ValueError(
+                "The given q value is too large for the given duration and "
+                f"sample rate. The maximum allowable value is {max_q}"
+            )
+
+        if interpolation_method not in ["bilinear", "bicubic", "spline"]:
+            raise ValueError(
+                "Interpolation method must be either 'bilinear', 'bicubic', "
+                f"or 'spline'; got {interpolation_method}"
+            )
+        self.interpolation_method = interpolation_method
+
         qprime = self.q / 11 ** (1 / 2.0)
         if self.frange[0] <= 0:  # set non-zero lower frequency
             self.frange[0] = 50 * self.q / (2 * torch.pi * duration)
         if math.isinf(self.frange[1]):  # set non-infinite upper frequency
             self.frange[1] = sample_rate / 2 / (1 + 1 / qprime)
+
         self.freqs = self.get_freqs()
         self.qtile_transforms = torch.nn.ModuleList(
             [
-                QTile(self.q, freq, self.duration, sample_rate, self.mismatch)
+                QTile(
+                    q=self.q,
+                    frequency=freq,
+                    duration=self.duration,
+                    sample_rate=sample_rate,
+                    mismatch=self.mismatch,
+                )
                 for freq in self.freqs
             ]
         )
         self.qtiles = None
 
+        if self.interpolation_method == "spline":
+            self._set_up_spline_interp()
+
+    def _set_up_spline_interp(self):
+        ntiles = [qtile.ntiles() for qtile in self.qtile_transforms]
+        # For efficiency, we'll stack all qtiles of the same length before
+        # interpolating, so we need to figure out which those are
+        unique_ntiles = sorted(list(set(ntiles)))
+        idx = torch.arange(len(ntiles))
+        self.stack_idx = [idx[Tensor(ntiles) == n] for n in unique_ntiles]
+
+        t_out = torch.arange(
+            0, self.duration, self.duration / self.spectrogram_shape[1]
+        )
+        self.qtile_interpolators = torch.nn.ModuleList(
+            [
+                SplineInterpolate(
+                    kx=3,
+                    x_in=torch.arange(0, self.duration, self.duration / tiles),
+                    y_in=torch.arange(len(idx)),
+                    x_out=t_out,
+                    y_out=torch.arange(len(idx)),
+                )
+                for tiles, idx in zip(unique_ntiles, self.stack_idx)
+            ]
+        )
+
+        t_in = t_out
+        f_in = self.freqs
+        f_out = torch.logspace(
+            math.log10(self.frange[0]),
+            math.log10(self.frange[-1]),
+            self.spectrogram_shape[0],
+        )
+
+        self.interpolator = SplineInterpolate(
+            kx=3,
+            ky=3,
+            x_in=t_in,
+            y_in=f_in,
+            x_out=t_out,
+            y_out=f_out,
+        )
+
     def get_freqs(self) -> Float[Tensor, " nfreq"]:
         """
         Calculate the frequencies that will be used in this transform.
@@ -220,7 +304,8 @@ class SingleQTransform(torch.nn.Module):
 
         freq_base = math.exp(2 / ((2 + self.q**2) ** (1 / 2.0)) * fstep)
         freqs = torch.Tensor([freq_base ** (i + 0.5) for i in range(nfreq)])
-        freqs = (minf * freqs // fstepmin) * fstepmin
+        # Cast freqs to float64 to avoid off-by-ones from rounding
+        freqs = (minf * freqs.double() // fstepmin) * fstepmin
         return torch.unique(freqs)
 
     def get_max_energy(
@@ -268,7 +353,11 @@ class SingleQTransform(torch.nn.Module):
         if dimension == "batch":
             return torch.max(max_across_ft, dim=-1).values
 
-    def compute_qtiles(self, X: TimeSeries1to3d, norm: str = "median") -> None:
+    def compute_qtiles(
+        self,
+        X: TimeSeries1to3d,
+        norm: str = "median",
+    ) -> None:
         """
         Take the FFT of the input timeseries and calculate the transform
         for each `QTile`
@@ -278,28 +367,40 @@ class SingleQTransform(torch.nn.Module):
         X[..., 1:] *= 2
         self.qtiles = [qtile(X, norm) for qtile in self.qtile_transforms]
 
-    def interpolate(self, num_f_bins: int, num_t_bins: int) -> TimeSeries3d:
-        """
-        Interpolate each `QTile` to the specified number of time and
-        frequency bins. Note that PyTorch does not have the same
-        interpolation methods that GWpy uses, and so the interpolated
-        spectrograms will be different even though the uninterpolated
-        values match. The `bicubic` interpolation method is used as
-        it seems to match GWpy most closely.
-        """
+    def interpolate(self) -> TimeSeries3d:
         if self.qtiles is None:
             raise RuntimeError(
                 "Q-tiles must first be computed with .compute_qtiles()"
             )
+        if self.interpolation_method == "spline":
+            qtiles = [
+                torch.stack([self.qtiles[i] for i in idx], dim=-2)
+                for idx in self.stack_idx
+            ]
+            time_interped = torch.cat(
+                [
+                    interpolator(qtile)
+                    for qtile, interpolator in zip(
+                        qtiles, self.qtile_interpolators
+                    )
+                ],
+                dim=-2,
+            )
+            return self.interpolator(time_interped)
+        num_f_bins, num_t_bins = self.spectrogram_shape
         resampled = [
             F.interpolate(
-                qtile[None], (qtile.shape[-2], num_t_bins), mode="bicubic"
+                qtile[None],
+                (qtile.shape[-2], num_t_bins),
+                mode=self.interpolation_method,
             )
             for qtile in self.qtiles
         ]
         resampled = torch.stack(resampled, dim=-2)
         resampled = F.interpolate(
-            resampled[0], (num_f_bins, num_t_bins), mode="bicubic"
+            resampled[0],
+            (num_f_bins, num_t_bins),
+            mode=self.interpolation_method,
         )
         return torch.squeeze(resampled)
 
@@ -307,7 +408,6 @@ class SingleQTransform(torch.nn.Module):
         self,
         X: TimeSeries1to3d,
         norm: str = "median",
-        spectrogram_shape: Optional[Tuple[int, int]] = None,
     ):
         """
         Compute the Q-tiles and interpolate
@@ -321,24 +421,15 @@ class SingleQTransform(torch.nn.Module):
                 three-dimensional, axes will be added during Q-tile
                 computation.
             norm:
-                The method of interpolation used by each QTile
-            spectrogram_shape:
-                The shape of the interpolated spectrogram, specified as
-                `(num_f_bins, num_t_bins)`. Because the
-                frequency spacing of the Q-tiles is in log-space, the frequency
-                interpolation is log-spaced as well. If not given, the shape
-                used to initialize the transform will be used.
+                The method of normalization used by each QTile
 
         Returns:
             The interpolated Q-transform for the batch of data. Output will
             have one more dimension than the input
         """
 
-        if spectrogram_shape is None:
-            spectrogram_shape = self.spectrogram_shape
-        num_f_bins, num_t_bins = spectrogram_shape
         self.compute_qtiles(X, norm)
-        return self.interpolate(num_f_bins, num_t_bins)
+        return self.interpolate()
 
 
 class QScan(torch.nn.Module):
@@ -376,14 +467,22 @@ class QScan(torch.nn.Module):
         spectrogram_shape: Tuple[int, int],
         qrange: List[float] = [4, 64],
         frange: List[float] = [0, torch.inf],
+        interpolation_method="bicubic",
         mismatch: float = 0.2,
     ) -> None:
         super().__init__()
         self.qrange = qrange
         self.mismatch = mismatch
-        self.qs = self.get_qs()
         self.frange = frange
         self.spectrogram_shape = spectrogram_shape
+        max_q = torch.pi * duration * sample_rate / 50 - 11 ** (0.5)
+        self.qs = self.get_qs()
+        if self.qs[-1] >= max_q:
+            warnings.warn(
+                "Some Q values exceed the maximum allowable Q value of "
+                f"{max_q}. The list of Q values to be tested in this "
+                "scan will be truncated to avoid those values."
+            )
 
         # Deliberately doing something different from GWpy here.
         # Their final frange is the intersection of the frange
@@ -397,9 +496,11 @@ class QScan(torch.nn.Module):
                     spectrogram_shape=spectrogram_shape,
                     q=q,
                     frange=self.frange.copy(),
+                    interpolation_method=interpolation_method,
                     mismatch=self.mismatch,
                 )
                 for q in self.qs
+                if q < max_q
             ]
         )
 
@@ -415,6 +516,7 @@ class QScan(torch.nn.Module):
             self.qrange[0] * math.exp(2 ** (1 / 2.0) * dq * (i + 0.5))
             for i in range(nplanes)
         ]
+
         return qs
 
     def forward(
@@ -422,7 +524,6 @@ class QScan(torch.nn.Module):
         X: TimeSeries1to3d,
         fsearch_range: List[float] = None,
         norm: str = "median",
-        spectrogram_shape: Optional[Tuple[int, int]] = None,
     ):
         """
         Compute the set of QTiles for each Q transform and determine which
@@ -442,12 +543,6 @@ class QScan(torch.nn.Module):
                 for the maximum energy
             norm:
                 The method of interpolation used by each QTile
-            spectrogram_shape:
-                The shape of the interpolated spectrogram, specified as
-                `(num_f_bins, num_t_bins)`. Because the
-                frequency spacing of the Q-tiles is in log-space, the frequency
-                interpolation is log-spaced as well. If not given, the shape
-                used to initialize the transform will be used.
 
         Returns:
             An interpolated Q-transform for the batch of data. Output will
@@ -463,7 +558,4 @@ class QScan(torch.nn.Module):
                 ]
             )
         )
-        if spectrogram_shape is None:
-            spectrogram_shape = self.spectrogram_shape
-        num_f_bins, num_t_bins = spectrogram_shape
-        return self.q_transforms[idx].interpolate(num_f_bins, num_t_bins)
+        return self.q_transforms[idx].interpolate()

{ml4gw-0.5.1 → ml4gw-0.6.0}/ml4gw/transforms/scaler.py

@@ -36,7 +36,9 @@ class ChannelWiseScaler(FittableTransform):
         self.register_buffer("mean", mean)
         self.register_buffer("std", std)
 
-    def fit(self, X: Float[Tensor, "... time"]) -> None:
+    def fit(
+        self, X: Float[Tensor, "... time"], std_reg: Optional[float] = 0.0
+    ) -> None:
         """Fit the scaling parameters to a timeseries
 
         Computes the channel-wise mean and standard deviation
@@ -59,7 +61,7 @@ class ChannelWiseScaler(FittableTransform):
                 "Can't fit channel wise mean and standard deviation "
                 "from tensor of shape {}".format(X.shape)
             )
-
+        std += std_reg * torch.ones_like(std)
         super().build(mean=mean, std=std)
 
     def forward(