ml4gw 0.4.0__py3-none-any.whl → 0.4.2__py3-none-any.whl

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

@@ -5,93 +5,109 @@ 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)
-
-
-class Cosine:
-    """
-    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
-    """
+    arg_constraints = {}
 
     def __init__(
-        self, low: float = -math.pi / 2, high: float = math.pi / 2
-    ) -> None:
+        self,
+        low: float = torch.as_tensor(-torch.pi / 2),
+        high: float = torch.as_tensor(torch.pi / 2),
+        validate_args=None,
+    ):
+        batch_shape = torch.Size()
+        super().__init__(batch_shape, validate_args=validate_args)
         self.low = low
-        self.norm = 1 / (math.sin(high) - math.sin(low))
+        self.norm = 1 / (torch.sin(high) - torch.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))
+    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 LogNormal:
+
+class Sine(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`.
+    Sine distribution based on
+    ``torch.distributions.TransformedDistribution``.
     """
 
     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 __call__(self, N: int) -> torch.Tensor:
-
-        u = self.mu + torch.randn(N) * self.sigma
-        x = torch.exp(u)
-
-        if self.low is not None:
-            x = torch.clip(x, self.low)
-        return x
-
-
-class LogUniform(Uniform):
+        self,
+        low: float = torch.as_tensor(0),
+        high: float = torch.as_tensor(torch.pi),
+        validate_args=None,
+    ):
+        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 uniform distribution
     """
 
-    def __init__(self, low: float, high: float) -> None:
-        super().__init__(math.log(low), math.log(high))
+    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 = super().__call__(N)
-        return 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)
 
-class PowerLaw:
+    def support(self):
+        if self.low is not None:
+            return dist.constraints.greater_than(self.low)
+
+
+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 +118,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 = torch.as_tensor(0.0),
+        validate_args=None,
+    ):
+        batch_shape = torch.Size()
+        super().__init__(batch_shape, validate_args=validate_args)
+        self.peak = 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/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

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/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):

ml4gw/waveforms/phenom_d.py

@@ -905,15 +905,15 @@ def phenom_d_inspiral_phase(Mf, mass_1, mass_2, eta, eta2, xi, chi1, chi2):
     sigma3 = sigma3Fit(eta, eta2, xi)
     sigma4 = sigma4Fit(eta, eta2, xi)
 
-    ins_phasing += (Mf.T * sigma1 / eta).T
-    ins_phasing += (Mf.T ** (4.0 / 3.0) * 0.75 * sigma2 / eta).T
-    ins_phasing += (Mf.T ** (5.0 / 3.0) * 0.6 * sigma3 / eta).T
-    ins_phasing += (Mf.T**2.0 * 0.5 * sigma4 / eta).T
+    ins_phasing += (Mf.mT * sigma1 / eta).mT
+    ins_phasing += (Mf.mT ** (4.0 / 3.0) * 0.75 * sigma2 / eta).mT
+    ins_phasing += (Mf.mT ** (5.0 / 3.0) * 0.6 * sigma3 / eta).mT
+    ins_phasing += (Mf.mT**2.0 * 0.5 * sigma4 / eta).mT
 
-    ins_Dphasing = (ins_Dphasing.T + sigma1 / eta).T
-    ins_Dphasing += (Mf.T ** (1.0 / 3.0) * sigma2 / eta).T
-    ins_Dphasing += (Mf.T ** (2.0 / 3.0) * sigma3 / eta).T
-    ins_Dphasing += (Mf.T * sigma4 / eta).T
+    ins_Dphasing = (ins_Dphasing.T + sigma1 / eta).mT
+    ins_Dphasing += (Mf.mT ** (1.0 / 3.0) * sigma2 / eta).mT
+    ins_Dphasing += (Mf.mT ** (2.0 / 3.0) * sigma3 / eta).mT
+    ins_Dphasing += (Mf.mT * sigma4 / eta).mT
 
     return ins_phasing, ins_Dphasing
 
@@ -925,16 +925,16 @@ def phenom_d_int_phase(Mf, eta, eta2, xi):
     # Merger phase
     # Leading beta0 is not added here
     # overall 1/eta is not multiplied
-    int_phasing = (Mf.T * beta1).T
-    int_phasing += (torch.log(Mf).T * beta2).T
-    int_phasing -= (Mf.T ** (-3.0) / 3.0 * beta3).T
+    int_phasing = (Mf.mT * beta1).mT
+    int_phasing += (torch.log(Mf).mT * beta2).mT
+    int_phasing -= (Mf.mT ** (-3.0) / 3.0 * beta3).mT
 
     # overall 1/eta is multiple in derivative of
     # intermediate phase
-    int_Dphasing = (Mf.T ** (-4.0) * beta3).T
-    int_Dphasing += (Mf.T ** (-1.0) * beta2).T
-    int_Dphasing = (int_Dphasing.T + beta1).T
-    int_Dphasing = (int_Dphasing.T / eta).T
+    int_Dphasing = (Mf.mT ** (-4.0) * beta3).mT
+    int_Dphasing += (Mf.mT ** (-1.0) * beta2).mT
+    int_Dphasing = (int_Dphasing.T + beta1).mT
+    int_Dphasing = (int_Dphasing.T / eta).mT
     return int_phasing, int_Dphasing
 
 
@@ -947,19 +947,21 @@ def phenom_d_mrd_phase(Mf, eta, eta2, chi1, chi2, xi):
 
     # merger ringdown
     fRD, fDM = fring_fdamp(eta, eta2, chi1, chi2)
-    f_minus_alpha5_fRD = (Mf.T - alpha5 * fRD).T
+    f_minus_alpha5_fRD = (Mf.t() - alpha5 * fRD).t()
 
     # Leading 1/eta is not multiplied at this stage
-    mrd_phasing = (Mf.T * alpha1).T
-    mrd_phasing -= (1 / Mf.T * alpha2).T
-    mrd_phasing += (4.0 / 3.0) * (Mf.T ** (3.0 / 4.0) * alpha3).T
-    mrd_phasing += (torch.atan(f_minus_alpha5_fRD.T / fDM) * alpha4).T
-
-    mrd_Dphasing = (alpha4 * fDM / (f_minus_alpha5_fRD.T**2 + fDM**2)).T
-    mrd_Dphasing += (Mf.T ** (-1.0 / 4.0) * alpha3).T
-    mrd_Dphasing += (Mf.T ** (-2.0) * alpha2).T
-    mrd_Dphasing = (mrd_Dphasing.T + alpha1).T
-    mrd_Dphasing = (mrd_Dphasing.T / eta).T
+    mrd_phasing = (Mf.t() * alpha1).t()
+    mrd_phasing -= (1 / Mf.t() * alpha2).t()
+    mrd_phasing += (4.0 / 3.0) * (Mf.t() ** (3.0 / 4.0) * alpha3).t()
+    mrd_phasing += (torch.atan(f_minus_alpha5_fRD.t() / fDM) * alpha4).t()
+
+    mrd_Dphasing = (
+        alpha4 * fDM / (f_minus_alpha5_fRD.t() ** 2 + fDM**2)
+    ).t()
+    mrd_Dphasing += (Mf.t() ** (-1.0 / 4.0) * alpha3).t()
+    mrd_Dphasing += (Mf.t() ** (-2.0) * alpha2).t()
+    mrd_Dphasing = (mrd_Dphasing.t() + alpha1).t()
+    mrd_Dphasing = (mrd_Dphasing.t() / eta).t()
 
     return mrd_phasing, mrd_Dphasing
 
@@ -984,23 +986,23 @@ def phenom_d_phase(Mf, mass_1, mass_2, eta, eta2, chi1, chi2, xi):
         PHI_fJoin_INS, eta, eta2, xi
     )
     C2Int = ins_Dphase_f1 - int_Dphase_f1
-    C1Int = ins_phase_f1 - (int_phase_f1.T / eta).T - C2Int * PHI_fJoin_INS
+    C1Int = ins_phase_f1 - (int_phase_f1.T / eta).mT - C2Int * PHI_fJoin_INS
     # C1 continuity at ringdown
-    fRDJoin = (0.5 * torch.ones_like(Mf).T * fRD).T
+    fRDJoin = (0.5 * torch.ones_like(Mf).mT * fRD).mT
     int_phase_rd, int_Dphase_rd = phenom_d_int_phase(fRDJoin, eta, eta2, xi)
     mrd_phase_rd, mrd_Dphase_rd = phenom_d_mrd_phase(
         fRDJoin, eta, eta2, chi1, chi2, xi
     )
-    PhiIntTempVal = (int_phase_rd.T / eta).T + C1Int + C2Int * fRDJoin
+    PhiIntTempVal = (int_phase_rd.T / eta).mT + C1Int + C2Int * fRDJoin
     # C2MRD = int_Dphase_rd - mrd_Dphase_rd
     C2MRD = C2Int + int_Dphase_rd - mrd_Dphase_rd
-    C1MRD = PhiIntTempVal - (mrd_phase_rd.T / eta).T - C2MRD * fRDJoin
+    C1MRD = PhiIntTempVal - (mrd_phase_rd.T / eta).mT - C2MRD * fRDJoin
 
-    int_phase = (int_phase.T / eta).T
+    int_phase = (int_phase.T / eta).mT
     int_phase += C1Int
     int_phase += Mf * C2Int
 
-    mrd_phase = (mrd_phase.T / eta).T
+    mrd_phase = (mrd_phase.T / eta).mT
     mrd_phase += C1MRD
     mrd_phase += Mf * C2MRD
 
@@ -1133,10 +1135,10 @@ def phenom_d_inspiral_amp(Mf, eta, eta2, Seta, xi, chi1, chi2, chi12, chi22):
         + Mf_five_third.T * prefactors_five_thirds
         + Mf_seven_third.T * prefactors_seven_thirds
         + MF_eight_third.T * prefactors_eight_thirds
-        + Mf.T * prefactors_one
+        + Mf.mT * prefactors_one
         + Mf_two.T * prefactors_two
         + Mf_three.T * prefactors_three
-    ).T
+    ).mT
 
     Damp = (
         (2.0 / 3.0) / Mf_one_third.T * prefactors_two_thirds
@@ -1145,9 +1147,9 @@ def phenom_d_inspiral_amp(Mf, eta, eta2, Seta, xi, chi1, chi2, chi12, chi22):
         + (7.0 / 3.0) * Mf_four_third.T * prefactors_seven_thirds
         + (8.0 / 3.0) * Mf_five_third.T * prefactors_eight_thirds
         + prefactors_one
-        + 2.0 * Mf.T * prefactors_two
+        + 2.0 * Mf.mT * prefactors_two
         + 3.0 * Mf_two.T * prefactors_three
-    ).T
+    ).mT
 
     return amp, Damp
 
@@ -1160,15 +1162,15 @@ def phenom_d_mrd_amp(Mf, eta, eta2, chi1, chi2, xi):
     gamma2 = gamma2_fun(eta, eta2, xi)
     gamma3 = gamma3_fun(eta, eta2, xi)
     fDMgamma3 = fDM * gamma3
-    pow2_fDMgamma3 = (torch.ones_like(Mf).T * fDMgamma3 * fDMgamma3).T
-    fminfRD = Mf - (torch.ones_like(Mf).T * fRD).T
-    exp_times_lorentzian = torch.exp(fminfRD.T * gamma2 / fDMgamma3).T
+    pow2_fDMgamma3 = (torch.ones_like(Mf).mT * fDMgamma3 * fDMgamma3).mT
+    fminfRD = Mf - (torch.ones_like(Mf).mT * fRD).mT
+    exp_times_lorentzian = torch.exp(fminfRD.mT * gamma2 / fDMgamma3).mT
     exp_times_lorentzian *= fminfRD**2 + pow2_fDMgamma3
 
-    amp = (1 / exp_times_lorentzian.T * gamma1 * gamma3 * fDM).T
-    Damp = (fminfRD.T * -2 * fDM * gamma1 * gamma3) / (
+    amp = (1 / exp_times_lorentzian.T * gamma1 * gamma3 * fDM).mT
+    Damp = (fminfRD.mT * -2 * fDM * gamma1 * gamma3) / (
         fminfRD * fminfRD + pow2_fDMgamma3
-    ).T - (gamma2 * gamma1)
+    ).mT - (gamma2 * gamma1)
     Damp = Damp.T / exp_times_lorentzian
     return amp, Damp
 
@@ -1184,7 +1186,7 @@ def phenom_d_int_amp(Mf, eta, eta2, Seta, chi1, chi2, chi12, chi22, xi):
     gamma3 = gamma3_fun(eta, eta2, xi)
 
     fpeak = fmaxCalc(fRD, fDM, gamma2, gamma3)
-    Mf3 = (torch.ones_like(Mf).T * fpeak).T
+    Mf3 = (torch.ones_like(Mf).mT * fpeak).mT
     dfx = 0.5 * (Mf3 - Mf1)
     Mf2 = Mf1 + dfx
 
@@ -1192,7 +1194,7 @@ def phenom_d_int_amp(Mf, eta, eta2, Seta, chi1, chi2, chi12, chi22, xi):
         Mf1, eta, eta2, Seta, xi, chi1, chi2, chi12, chi22
     )
     v3, d2 = phenom_d_mrd_amp(Mf3, eta, eta2, chi1, chi2, xi)
-    v2 = (torch.ones_like(Mf).T * AmpIntColFitCoeff(eta, eta2, xi)).T
+    v2 = (torch.ones_like(Mf).mT * AmpIntColFitCoeff(eta, eta2, xi)).mT
 
     delta_0, delta_1, delta_2, delta_3, delta_4 = delta_values(
         f1=Mf1, f2=Mf2, f3=Mf3, v1=v1, v2=v2, v3=v3, d1=d1, d2=d2
@@ -1225,7 +1227,7 @@ def phenom_d_amp(
     fRD, fDM = fring_fdamp(eta, eta2, chi1, chi2)
     Mf_peak = fmaxCalc(fRD, fDM, gamma2, gamma3)
     # Geometric peak and joining frequencies
-    Mf_peak = (torch.ones_like(Mf).T * Mf_peak).T
+    Mf_peak = (torch.ones_like(Mf).mT * Mf_peak).mT
     Mf_join_ins = 0.014 * torch.ones_like(Mf)
 
     # construct full IMR Amp
@@ -1290,9 +1292,9 @@ def phenom_d_htilde(
         Mf_ref, mass_1, mass_2, eta, eta2, chi1, chi2, xi
     )
 
-    Psi = (Psi.T - 2 * phic).T
+    Psi = (Psi.T - 2 * phic).mT
     Psi -= Psi_ref
-    Psi -= ((Mf - Mf_ref).T * t0).T
+    Psi -= ((Mf - Mf_ref).mT * t0).mT
 
     amp, _ = phenom_d_amp(
         Mf,
@@ -1353,7 +1355,7 @@ def IMRPhenomD(
         f, chirp_mass, mass_ratio, chi1, chi2, distance, phic, f_ref
     )
 
-    hp = (htilde.T * pfac).T
-    hc = -1j * (htilde.T * cfac).T
+    hp = (htilde.mT * pfac).mT
+    hc = -1j * (htilde.mT * cfac).mT
 
     return hp, hc

{ml4gw-0.4.0.dist-info → ml4gw-0.4.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ml4gw
-Version: 0.4.0
+Version: 0.4.2
 Summary: Tools for training torch models on gravitational wave data
 Author: Alec Gunny
 Author-email: alec.gunny@ligo.org
@@ -10,10 +10,8 @@ Classifier: Programming Language :: Python :: 3.8
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
-Requires-Dist: torch (>=1.10,<2.0) ; python_version >= "3.8" and python_version < "3.11"
-Requires-Dist: torch (>=2.0,<3.0) ; python_version >= "3.11"
-Requires-Dist: torchaudio (>=0.13,<0.14) ; python_version >= "3.8" and python_version < "3.11"
-Requires-Dist: torchaudio (>=2.0,<3.0) ; python_version >= "3.11"
+Requires-Dist: torch (>=2.0,<3.0)
+Requires-Dist: torchaudio (>=2.0,<3.0)
 Requires-Dist: torchtyping (>=0.1,<0.2)
 Description-Content-Type: text/markdown
 

{ml4gw-0.4.0.dist-info → ml4gw-0.4.2.dist-info}/RECORD

@@ -3,8 +3,8 @@ ml4gw/augmentations.py,sha256=UEwNxdjzIQUqUgtAdrstptTZ73ay5sLc3imQY1DVUqs,1027
 ml4gw/dataloading/__init__.py,sha256=EHBBqU7y2-Np5iQ_xyufxamUEM1pPEquqFo7oaJnaJE,149
 ml4gw/dataloading/chunked_dataset.py,sha256=jy-y5xhMJqRZIA_pjrc4QHhqpAcpM9aJi2omT24riXY,5195
 ml4gw/dataloading/hdf5_dataset.py,sha256=D6cWBtF5_nox89SO4M9so2GDhkfkhdkc0EUpPhyvEyE,6643
-ml4gw/dataloading/in_memory_dataset.py,sha256=9P0mk0PPno6ErRrxX6nhcYuwTJWAUAp7VSls2DBFRwo,9801
-ml4gw/distributions.py,sha256=m38ynrUy0QoCpqkyLHSzUZe72tA02se5UBXuXQhN4bQ,3428
+ml4gw/dataloading/in_memory_dataset.py,sha256=8bHOB7GreoLWcb_IYEQd2-BXfaYCO1EonD6MHFVPAzA,9429
+ml4gw/distributions.py,sha256=l_GBHCpWsAASawvACBy2vdZA5KTabcH7gnoRClstUjA,4838
 ml4gw/gw.py,sha256=RlG8Vj9GCcv0wpKWz9XlbRattls1A-FCCN1RiIhSE_w,17914
 ml4gw/nn/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ml4gw/nn/autoencoder/__init__.py,sha256=ZaT1XhJTHpMuPQqu5E__Jezeh9uwtjcXlT7IZ18byq4,161
@@ -20,8 +20,9 @@ ml4gw/nn/streaming/__init__.py,sha256=zgjGR2L8t0txXLnil9ceZT0tM8Y2FC8yPxqIKYH0o1
 ml4gw/nn/streaming/online_average.py,sha256=T-wWw7eEufbUVPRNnLAXIq0cedAyJWEE9tdZ6CTi3cs,4561
 ml4gw/nn/streaming/snapshotter.py,sha256=-l_YsWby7ZnEzGIAlLAV2mtR0daLMtLCxovtt4OI3Z0,4432
 ml4gw/spectral.py,sha256=5GfKAV_1vw5yyzTD2u_myjT5jIlAyAHDX6TXj9ynL_o,19021
-ml4gw/transforms/__init__.py,sha256=t6ZJcq23apqDKhLGM-U5l_bqxJcXFj3riY6cTGY47Gc,314
+ml4gw/transforms/__init__.py,sha256=24pdP_hIg1wfrtZxxRBPhcEXsCbvVKtNKp7JL8SEogE,362
 ml4gw/transforms/pearson.py,sha256=bJ77lO4wBY6y1R1aESN_bcUEMbc55hWCIaCBdbIj4CY,3133
+ml4gw/transforms/qtransform.py,sha256=hvCzdGROLoW1nJYR_ZZWDnafJpX4kD1os3CZ2jQJ7IU,17328
 ml4gw/transforms/scaler.py,sha256=5VGov0M80NZostRzccViC3HNftx4ZVu0kOKTDmiLrR4,2327
 ml4gw/transforms/snr_rescaler.py,sha256=ocYr6UjpHW7t5TvruV7fyY8KuuDfGOJyvxEulmiFA6o,2275
 ml4gw/transforms/spectral.py,sha256=Vba9199z_ZaxsHWxdpgHB3U216rmGoSyehtvM3R9Z7A,3771
@@ -33,11 +34,11 @@ ml4gw/types.py,sha256=XbxunX8zRF95Fp1mZ9jEbixb63bwDQMoayRMMxT9Lzo,429
 ml4gw/utils/interferometer.py,sha256=w_0WkboCJZMKAg-4lhiNGOOkNogAghpT96I0TE5aJ1g,1519
 ml4gw/utils/slicing.py,sha256=Cbwcpk_0hsfN4zczFVM2YbDRjeirA7jFvApM4Jy0U8s,13535
 ml4gw/waveforms/__init__.py,sha256=zjqOKNY4z1A5iPhWTxyhnkLh2robB-obPTtaK-pDUoU,104
-ml4gw/waveforms/generator.py,sha256=4Z6vUEuI84t__3t0DDnXlOyB8R96ynf8xFvtwCGu9JA,1057
-ml4gw/waveforms/phenom_d.py,sha256=pxHk7paW5709Ak29m_DYeQ8kiMLC8wrUnM13flUU36o,38419
+ml4gw/waveforms/generator.py,sha256=HYTAbih5y-i4v5iV4twTo2DPSADUrrsuz2m9WDqJoH4,1067
+ml4gw/waveforms/phenom_d.py,sha256=OVWkGKmBCwcE6xFuaJri4F9bh-Pmkwt_lgZv7wNEtsU,38535
 ml4gw/waveforms/phenom_d_data.py,sha256=WA1FBxUp9fo1IQaV_OLJ_5g5gI166mY1FtG9n25he9U,53447
 ml4gw/waveforms/sine_gaussian.py,sha256=WZ6KiVEFSjB9Tv5otJbvI_Yr3341th1Noec_LB9kPOE,3577
 ml4gw/waveforms/taylorf2.py,sha256=x3drvKUMarWI9xHUzMRQhVp1Hh7X-j5WC2bdsbEiVfk,8482
-ml4gw-0.4.0.dist-info/METADATA,sha256=SyBcghsG_wj1TSYjmCnklojeCAH0OQ3uv-DXOoiiuug,5944
-ml4gw-0.4.0.dist-info/WHEEL,sha256=FMvqSimYX_P7y0a7UY-_Mc83r5zkBZsCYPm7Lr0Bsq4,88
-ml4gw-0.4.0.dist-info/RECORD,,
+ml4gw-0.4.2.dist-info/METADATA,sha256=Zbam_nfw6oAeAzBSSBbJpJQK2egAiMmbWFt6eidTNIQ,5706
+ml4gw-0.4.2.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+ml4gw-0.4.2.dist-info/RECORD,,

{ml4gw-0.4.0.dist-info → ml4gw-0.4.2.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: poetry-core 1.8.1
+Generator: poetry-core 1.9.0
 Root-Is-Purelib: true
 Tag: py3-none-any