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

ml4gw/transforms/qtransform.py

@@ -0,0 +1,463 @@
+import math
+from typing import List, Optional, Tuple
+
+import torch
+import torch.nn.functional as F
+
+"""
+All based on https://github.com/gwpy/gwpy/blob/v3.0.8/gwpy/signal/qtransform.py
+The methods, names, and descriptions come almost entirely from GWpy.
+This code allows the Q-transform to be performed on batches of multi-channel
+input on GPU.
+"""
+
+
+class QTile(torch.nn.Module):
+    """
+    Compute the row of Q-tiles for a single Q value and a single
+    frequency for a batch of multi-channel frequency series data.
+    Should really be called `QRow`, but I want to match GWpy.
+    Input data should have three dimensions or fewer.
+    If fewer, dimensions will be added until the input is
+    three-dimensional.
+
+    Args:
+        q:
+            The Q value to use in computing the Q tile
+        frequency:
+            The frequency for which to compute the Q tile in Hz
+        duration:
+            The length of time in seconds that the input frequency
+            series represents
+        sample_rate:
+            The sample rate of the original time series in Hz
+        mismatch:
+            The maximum fractional mismatch between neighboring tiles
+
+
+    """
+
+    def __init__(
+        self,
+        q: float,
+        frequency: float,
+        duration: float,
+        sample_rate: float,
+        mismatch: float,
+    ):
+        super().__init__()
+        self.mismatch = mismatch
+        self.q = q
+        self.deltam = torch.tensor(2 * (self.mismatch / 3.0) ** (1 / 2.0))
+        self.qprime = self.q / 11 ** (1 / 2.0)
+        self.frequency = frequency
+        self.duration = duration
+        self.sample_rate = sample_rate
+
+        self.windowsize = (
+            2 * int(self.frequency / self.qprime * self.duration) + 1
+        )
+        pad = self.ntiles() - self.windowsize
+        padding = torch.Tensor((int((pad - 1) / 2.0), int((pad + 1) / 2.0)))
+        self.register_buffer("padding", padding)
+        self.register_buffer("indices", self.get_data_indices())
+        self.register_buffer("window", self.get_window())
+
+    def ntiles(self):
+        """
+        Number of tiles in this frequency row
+        """
+        tcum_mismatch = self.duration * 2 * torch.pi * self.frequency / self.q
+        return int(2 ** torch.ceil(torch.log2(tcum_mismatch / self.deltam)))
+
+    def _get_indices(self):
+        half = int((self.windowsize - 1) / 2)
+        return torch.arange(-half, half + 1)
+
+    def get_window(self):
+        """
+        Generate the bi-square window for this row
+        """
+        wfrequencies = self._get_indices() / self.duration
+        xfrequencies = wfrequencies * self.qprime / self.frequency
+        norm = (
+            self.ntiles()
+            / (self.duration * self.sample_rate)
+            * (315 * self.qprime / (128 * self.frequency)) ** (1 / 2.0)
+        )
+        return torch.Tensor((1 - xfrequencies**2) ** 2 * norm)
+
+    def get_data_indices(self):
+        """
+        Get the index array of relevant frequencies for this row
+        """
+        return torch.round(
+            self._get_indices() + 1 + self.frequency * self.duration,
+        ).type(torch.long)
+
+    def forward(self, fseries: torch.Tensor, norm: str = "median"):
+        """
+        Compute the transform for this row
+
+        Args:
+            fseries:
+                Frequency series of data. Should correspond to data with
+                the duration and sample rate used to initialize this object.
+                Expected input shape is `(B, C, F)`, where F is the number
+                of samples, C is the number of channels, and B is the number
+                of batches. If less than three-dimensional, axes will be
+                added.
+            norm:
+                The method of normalization. Options are "median", "mean", or
+                `None`.
+
+        Returns:
+            The row of Q-tiles for the given Q and frequency. Output is
+            three-dimensional: `(B, C, T)`
+        """
+        if len(fseries.shape) > 3:
+            raise ValueError("Input data has more than 3 dimensions")
+
+        while len(fseries.shape) < 3:
+            fseries = fseries[None]
+
+        windowed = fseries[..., self.indices] * self.window
+        left, right = self.padding
+        padded = F.pad(windowed, (int(left), int(right)), mode="constant")
+        wenergy = torch.fft.ifftshift(padded, dim=-1)
+
+        tdenergy = torch.fft.ifft(wenergy)
+        energy = tdenergy.real**2.0 + tdenergy.imag**2.0
+        if norm:
+            norm = norm.lower() if isinstance(norm, str) else norm
+            if norm == "median":
+                medians = torch.quantile(energy, q=0.5, dim=-1, keepdim=True)
+                energy /= medians
+            elif norm == "mean":
+                means = torch.mean(energy, dim=-1, keepdim=True)
+                energy /= means
+            else:
+                raise ValueError("Invalid normalisation %r" % norm)
+            return energy.type(torch.float32)
+        return energy
+
+
+class SingleQTransform(torch.nn.Module):
+    """
+    Compute the Q-transform for a single Q value for a batch of
+    multi-channel time series data. Input data should have
+    three dimensions or fewer.
+
+    Args:
+        duration:
+            Length of the time series data in seconds
+        sample_rate:
+            Sample rate of the data in Hz
+        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.
+        q:
+            The Q value to use for the Q transform
+        frange:
+            The lower and upper frequency limit to consider for
+            the transform. If unspecified, default values will
+            be chosen based on q, sample_rate, and duration
+        mismatch:
+            The maximum fractional mismatch between neighboring tiles
+    """
+
+    def __init__(
+        self,
+        duration: float,
+        sample_rate: float,
+        spectrogram_shape: Tuple[int, int],
+        q: float = 12,
+        frange: List[float] = [0, torch.inf],
+        mismatch: float = 0.2,
+    ):
+        super().__init__()
+        self.q = q
+        self.spectrogram_shape = spectrogram_shape
+        self.frange = frange
+        self.duration = duration
+        self.mismatch = mismatch
+
+        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)
+                for freq in self.freqs
+            ]
+        )
+        self.qtiles = None
+
+    def get_freqs(self):
+        """
+        Calculate the frequencies that will be used in this transform.
+        For each frequency, a `QTile` is created.
+        """
+        minf, maxf = self.frange
+        fcum_mismatch = (
+            math.log(maxf / minf) * (2 + self.q**2) ** (1 / 2.0) / 2.0
+        )
+        deltam = 2 * (self.mismatch / 3.0) ** (1 / 2.0)
+        nfreq = int(max(1, math.ceil(fcum_mismatch / deltam)))
+        fstep = fcum_mismatch / nfreq
+        fstepmin = 1 / self.duration
+
+        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
+        return torch.unique(freqs)
+
+    def get_max_energy(
+        self, fsearch_range: List[float] = None, dimension: str = "both"
+    ):
+        """
+        Gets the maximum energy value among the QTiles. The maximum can
+        be computed across all batches and channels, across all channels,
+        across all batches, or individually for each channel/batch
+        combination. This could be useful for allowing the use of different
+        Q values for different channels and batches, but the slicing would
+        be slow, so this isn't used yet.
+
+        Optionally, a pair of frequency values can be specified for
+        `fsearch_range` to restrict the frequencies in which the maximum
+        energy value is sought.
+        """
+        allowed_dimensions = ["both", "neither", "channel", "batch"]
+        if dimension not in allowed_dimensions:
+            raise ValueError(f"Dimension must be one of {allowed_dimensions}")
+
+        if self.qtiles is None:
+            raise RuntimeError(
+                "Q-tiles must first be computed with .compute_qtiles()"
+            )
+
+        if fsearch_range is not None:
+            start = min(torch.argwhere(self.freqs > fsearch_range[0]))
+            stop = min(torch.argwhere(self.freqs > fsearch_range[1]))
+            qtiles = self.qtiles[start:stop]
+        else:
+            qtiles = self.qtiles
+
+        if dimension == "both":
+            return max([torch.max(qtile) for qtile in qtiles])
+
+        max_across_t = [torch.max(qtile, dim=-1).values for qtile in qtiles]
+        max_across_t = torch.stack(max_across_t, dim=-1)
+        max_across_ft = torch.max(max_across_t, dim=-1).values
+
+        if dimension == "neither":
+            return max_across_ft
+        if dimension == "channel":
+            return torch.max(max_across_ft, dim=-2).values
+        if dimension == "batch":
+            return torch.max(max_across_ft, dim=-1).values
+
+    def compute_qtiles(self, X: torch.Tensor, norm: str = "median"):
+        """
+        Take the FFT of the input timeseries and calculate the transform
+        for each `QTile`
+        """
+        # Computing the FFT with the same normalization and scaling as GWpy
+        X = torch.fft.rfft(X, norm="forward")
+        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):
+        """
+        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.
+        """
+        if self.qtiles is None:
+            raise RuntimeError(
+                "Q-tiles must first be computed with .compute_qtiles()"
+            )
+        resampled = [
+            F.interpolate(
+                qtile[None], (qtile.shape[-2], num_t_bins), mode="bicubic"
+            )
+            for qtile in self.qtiles
+        ]
+        resampled = torch.stack(resampled, dim=-2)
+        resampled = F.interpolate(
+            resampled[0], (num_f_bins, num_t_bins), mode="bicubic"
+        )
+        return torch.squeeze(resampled)
+
+    def forward(
+        self,
+        X: torch.Tensor,
+        norm: str = "median",
+        spectrogram_shape: Optional[Tuple[int, int]] = None,
+    ):
+        """
+        Compute the Q-tiles and interpolate
+
+        Args:
+            X:
+                Time series of data. Should have the duration and sample rate
+                used to initialize this object. Expected input shape is
+                `(B, C, T)`, where T is the number of samples, C is the number
+                of channels, and B is the number of batches. If less than
+                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.
+
+        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)
+
+
+class QScan(torch.nn.Module):
+    """
+    Calculate the Q-transform of a batch of multi-channel
+    time series data for a range of Q values and return
+    the interpolated Q-transform with the highest energy.
+
+    Args:
+        duration:
+            Length of the time series data in seconds
+        sample_rate:
+            Sample rate of the data in Hz
+        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.
+        qrange:
+            The lower and upper values of Q to consider. The
+            actual values of Q used for the transforms are
+            determined by the `get_qs` method
+        frange:
+            The lower and upper frequency limit to consider for
+            the transform. If unspecified, default values will
+            be chosen based on q, sample_rate, and duration
+        mismatch:
+            The maximum fractional mismatch between neighboring tiles
+    """
+
+    def __init__(
+        self,
+        duration: float,
+        sample_rate: float,
+        spectrogram_shape: Tuple[int, int],
+        qrange: List[float] = [4, 64],
+        frange: List[float] = [0, torch.inf],
+        mismatch: float = 0.2,
+    ):
+        super().__init__()
+        self.qrange = qrange
+        self.mismatch = mismatch
+        self.qs = self.get_qs()
+        self.frange = frange
+        self.spectrogram_shape = spectrogram_shape
+
+        # Deliberately doing something different from GWpy here.
+        # Their final frange is the intersection of the frange
+        # from each q. This implementation uses the frange of
+        # the chosen q.
+        self.q_transforms = torch.nn.ModuleList(
+            [
+                SingleQTransform(
+                    duration=duration,
+                    sample_rate=sample_rate,
+                    spectrogram_shape=spectrogram_shape,
+                    q=q,
+                    frange=self.frange.copy(),
+                    mismatch=self.mismatch,
+                )
+                for q in self.qs
+            ]
+        )
+
+    def get_qs(self):
+        """
+        Determine the values of Q to try for the set of Q-transforms
+        """
+        deltam = 2 * (self.mismatch / 3.0) ** (1 / 2.0)
+        cumum = math.log(self.qrange[1] / self.qrange[0]) / 2 ** (1 / 2.0)
+        nplanes = int(max(math.ceil(cumum / deltam), 1))
+        dq = cumum / nplanes
+        qs = [
+            self.qrange[0] * math.exp(2 ** (1 / 2.0) * dq * (i + 0.5))
+            for i in range(nplanes)
+        ]
+        return qs
+
+    def forward(
+        self,
+        X: torch.Tensor,
+        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
+        has the highest energy value. Interpolate and return the
+        corresponding set of tiles.
+
+        Args:
+            X:
+                Time series of data. Should have the duration and sample rate
+                used to initialize this object. Expected input shape is
+                `(B, C, T)`, where T is the number of samples, C is the number
+                of channels, and B is the number of batches. If less than
+                three-dimensional, axes will be added during Q-tile
+                computation.
+            fsearch_range:
+                The lower and upper frequency values within which to search
+                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
+            have one more dimension than the input
+        """
+        for transform in self.q_transforms:
+            transform.compute_qtiles(X, norm)
+        idx = torch.argmax(
+            torch.Tensor(
+                [
+                    transform.get_max_energy(fsearch_range=fsearch_range)
+                    for transform in self.q_transforms
+                ]
+            )
+        )
+        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)

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

ml4gw/waveforms/generator.py

@@ -12,7 +12,7 @@ class ParameterSampler(torch.nn.Module):
         self,
         N: int,
     ):
-        return {k: v(N) for k, v in self.parameters.items()}
+        return {k: v.sample((N,)) for k, v in self.parameters.items()}
 
 
 class WaveformGenerator(torch.nn.Module):