ml4gw 0.6.3__py3-none-any.whl → 0.7.1__py3-none-any.whl

ml4gw/__init__.py

@@ -0,0 +1 @@
+from .constants import *

ml4gw/dataloading/chunked_dataset.py

@@ -2,7 +2,7 @@ from collections.abc import Iterable
 
 import torch
 
-from ml4gw.types import WaveformTensor
+from ..types import WaveformTensor
 
 
 class ChunkedTimeSeriesDataset(torch.utils.data.IterableDataset):

ml4gw/dataloading/hdf5_dataset.py

@@ -1,11 +1,11 @@
 import warnings
-from typing import Sequence, Union
+from typing import Optional, Sequence, Union
 
 import h5py
 import numpy as np
 import torch
 
-from ml4gw.types import WaveformTensor
+from ..types import WaveformTensor
 
 
 class ContiguousHdf5Warning(Warning):
@@ -50,6 +50,13 @@ class Hdf5TimeSeriesDataset(torch.utils.data.IterableDataset):
             channel. The latter setting limits the amount of
             entropy in the effective dataset, but can provide
             over 2x improvement in total throughput.
+        num_files_per_batch:
+            The number of unique files from which to sample
+            batch elements each epoch. If left as `None`,
+            will use all available files. Useful when reading
+            from many files is bottlenecking dataloading.
+
+
     """
 
     def __init__(
@@ -60,6 +67,7 @@ class Hdf5TimeSeriesDataset(torch.utils.data.IterableDataset):
         batch_size: int,
         batches_per_epoch: int,
         coincident: Union[bool, str],
+        num_files_per_batch: Optional[int] = None,
     ) -> None:
         if not isinstance(coincident, bool) and coincident != "files":
             raise ValueError(
@@ -67,13 +75,21 @@ class Hdf5TimeSeriesDataset(torch.utils.data.IterableDataset):
                 "got unrecognized value {}".format(coincident)
             )
 
-        self.fnames = fnames
+        self.fnames = np.array(fnames)
         self.channels = channels
         self.num_channels = len(channels)
         self.kernel_size = kernel_size
         self.batch_size = batch_size
         self.batches_per_epoch = batches_per_epoch
         self.coincident = coincident
+        self.num_files_per_batch = (
+            len(fnames) if num_files_per_batch is None else num_files_per_batch
+        )
+        if self.num_files_per_batch > len(fnames):
+            raise ValueError(
+                f"Number of files per batch ({self.num_files_per_batch}) "
+                f"cannot exceed number of files ({len(fnames)}) "
+            )
 
         self.sizes = {}
         for fname in self.fnames:
@@ -85,13 +101,14 @@ class Hdf5TimeSeriesDataset(torch.utils.data.IterableDataset):
                         "without using chunked storage. This can have "
                         "severe performance impacts at data loading time. "
                         "If you need faster loading, try re-generating "
-                        "your datset with chunked storage turned on.".format(
+                        "your dataset with chunked storage turned on.".format(
                             fname
                         ),
                         category=ContiguousHdf5Warning,
                     )
 
                 self.sizes[fname] = len(dset)
+
         total = sum(self.sizes.values())
         self.probs = np.array([i / total for i in self.sizes.values()])
 
@@ -99,9 +116,22 @@ class Hdf5TimeSeriesDataset(torch.utils.data.IterableDataset):
         return self.batches_per_epoch
 
     def sample_fnames(self, size) -> np.ndarray:
-        return np.random.choice(
-            self.fnames,
+        # first, randomly select `self.num_files_per_batch`
+        # file indices based on their probabilities
+        fname_indices = np.arange(len(self.fnames))
+        fname_indices = np.random.choice(
+            fname_indices,
             p=self.probs,
+            size=(self.num_files_per_batch),
+            replace=False,
+        )
+        # now renormalize the probabilities, and sample
+        # the requested size from this subset of files
+        probs = self.probs[fname_indices]
+        probs /= probs.sum()
+        return np.random.choice(
+            self.fnames[fname_indices],
+            p=probs,
             size=size,
             replace=True,
         )

ml4gw/dataloading/in_memory_dataset.py

@@ -5,7 +5,7 @@ import torch
 from jaxtyping import Float
 from torch import Tensor
 
-from ml4gw.utils.slicing import slice_kernels
+from ..utils.slicing import slice_kernels
 
 
 class InMemoryDataset(torch.utils.data.IterableDataset):

ml4gw/gw.py

@@ -16,8 +16,10 @@ import torch
 from jaxtyping import Float
 from torch import Tensor
 
-from ml4gw.constants import C
-from ml4gw.types import (
+from ml4gw.utils.interferometer import InterferometerGeometry
+
+from .constants import C
+from .types import (
     BatchTensor,
     NetworkDetectorTensors,
     NetworkVertices,
@@ -26,7 +28,6 @@ from ml4gw.types import (
     VectorGeometry,
     WaveformTensor,
 )
-from ml4gw.utils.interferometer import InterferometerGeometry
 
 
 def outer(x: VectorGeometry, y: VectorGeometry) -> TensorGeometry:

ml4gw/nn/autoencoder/base.py

@@ -4,7 +4,7 @@ from typing import Optional, Tuple, Union
 import torch
 from torch import Tensor
 
-from ml4gw.nn.autoencoder.skip_connection import SkipConnection
+from .skip_connection import SkipConnection
 
 
 class Autoencoder(torch.nn.Module):

ml4gw/nn/autoencoder/convolutional.py

@@ -4,9 +4,9 @@ from typing import Optional
 import torch
 from torch import Tensor
 
-from ml4gw.nn.autoencoder.base import Autoencoder
-from ml4gw.nn.autoencoder.skip_connection import SkipConnection
-from ml4gw.nn.autoencoder.utils import match_size
+from .base import Autoencoder
+from .skip_connection import SkipConnection
+from .utils import match_size
 
 Module = Callable[[...], torch.nn.Module]
 

ml4gw/nn/autoencoder/skip_connection.py

@@ -1,7 +1,7 @@
 import torch
 from torch import Tensor
 
-from ml4gw.nn.autoencoder.utils import match_size
+from .utils import match_size
 
 
 class SkipConnection(torch.nn.Module):

ml4gw/nn/resnet/resnet_1d.py

@@ -13,7 +13,7 @@ import torch
 import torch.nn as nn
 from torch import Tensor
 
-from ml4gw.nn.norm import GroupNorm1DGetter, NormLayer
+from ..norm import GroupNorm1DGetter, NormLayer
 
 
 def convN(

ml4gw/nn/resnet/resnet_2d.py

@@ -10,7 +10,7 @@ import torch
 import torch.nn as nn
 from torch import Tensor
 
-from ml4gw.nn.norm import GroupNorm2DGetter, NormLayer
+from ..norm import GroupNorm2DGetter, NormLayer
 
 
 def convN(

ml4gw/nn/streaming/online_average.py

@@ -4,7 +4,7 @@ import torch
 from jaxtyping import Float
 from torch import Tensor
 
-from ml4gw.utils.slicing import unfold_windows
+from ...utils.slicing import unfold_windows
 
 
 class OnlineAverager(torch.nn.Module):

ml4gw/nn/streaming/snapshotter.py

@@ -4,7 +4,7 @@ import torch
 from jaxtyping import Float
 from torch import Tensor
 
-from ml4gw.utils.slicing import unfold_windows
+from ...utils.slicing import unfold_windows
 
 
 class Snapshotter(torch.nn.Module):

ml4gw/spectral.py

@@ -15,7 +15,7 @@ import torch
 from jaxtyping import Float
 from torch import Tensor
 
-from ml4gw.types import (
+from .types import (
     FrequencySeries1to3d,
     PSDTensor,
     TimeSeries1to3d,
@@ -343,6 +343,7 @@ def truncate_inverse_power_spectrum(
     fduration: Union[Float[Tensor, " time"], float],
     sample_rate: float,
     highpass: Optional[float] = None,
+    lowpass: Optional[float] = None,
 ) -> PSDTensor:
     """
     Truncate the length of the time domain response
@@ -375,6 +376,10 @@ def truncate_inverse_power_spectrum(
             If specified, will zero out the frequency response
             of all frequencies below this value in Hz. If left
             as `None`, no highpass filtering will be applied.
+        lowpass:
+            If specified, will zero out the frequency response
+            of all frequencies above this value in Hz. If left
+            as `None`, no lowpass filtering will be applied.
     Returns:
         The PSD with its time domain response truncated
             to `fduration` and any highpassed frequencies
@@ -388,12 +393,15 @@ def truncate_inverse_power_spectrum(
     # impulse response function
     inv_asd = 1 / psd**0.5
 
-    # zero our leading frequencies if we want the
-    # filter to perform highpass filtering
+    # zero out frequencies if we want the filter
+    # to perform highpass/lowpass filtering
+    df = sample_rate / N
     if highpass is not None:
-        df = sample_rate / N
         idx = int(highpass / df)
         inv_asd[:, :, :idx] = 0
+    if lowpass is not None:
+        idx = int(lowpass / df)
+        inv_asd[:, :, idx:] = 0
 
     if inv_asd.size(-1) % 2:
         inv_asd[:, :, -1] = 0
@@ -455,12 +463,13 @@ def whiten(
     fduration: Union[Float[Tensor, " time"], float],
     sample_rate: float,
     highpass: Optional[float] = None,
+    lowpass: Optional[float] = None,
 ) -> WaveformTensor:
     """
     Whiten a batch of timeseries using the specified
     background one-sided power spectral densities (PSDs),
     modified to have the desired time domain response length
-    `fduration` and possibly to highpass filter.
+    `fduration` and possibly to highpass/lowpass filter.
 
     Args:
         X:
@@ -493,6 +502,11 @@ def whiten(
             the data, setting the frequency response in the
             whitening filter to 0. If left as `None`, no
             highpass filtering will be applied.
+        lowpass:
+            The frequency in Hz at which to lowpass filter
+            the data, setting the frequency response in the
+            whitening filter to 0. If left as `None`, no
+            lowpass filtering will be applied.
     Returns:
         Batch of whitened multichannel timeseries with
             `fduration / 2` seconds trimmed from each side.
@@ -529,7 +543,11 @@ def whiten(
     # truncate it to have the desired
     # time domain response length
     psd = truncate_inverse_power_spectrum(
-        psd, fduration, sample_rate, highpass
+        psd,
+        fduration,
+        sample_rate,
+        highpass,
+        lowpass,
     )
 
     return normalize_by_psd(X, psd, sample_rate, pad)

ml4gw/transforms/__init__.py

@@ -1,3 +1,4 @@
+from .iirfilter import IIRFilter
 from .pearson import ShiftedPearsonCorrelation
 from .qtransform import QScan, SingleQTransform
 from .scaler import ChannelWiseScaler

ml4gw/transforms/iirfilter.py

@@ -0,0 +1,100 @@
+from typing import Union
+
+import torch
+from scipy.signal import iirfilter
+from torchaudio.functional import filtfilt
+
+
+class IIRFilter(torch.nn.Module):
+    r"""
+    IIR digital and analog filter design given order and critical points.
+    Design an Nth-order digital or analog filter and apply it to a signal.
+    Uses SciPy's `iirfilter` function to create the filter coefficients.
+    https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.iirfilter.html # noqa E501
+
+    The forward call of this module accepts a batch tensor of shape
+    (n_waveforms, n_samples) and returns the filtered waveforms.
+
+    Args:
+        N:
+            The order of the filter.
+        Wn:
+            A scalar or length-2 sequence giving the critical frequencies.
+            For digital filters, Wn are in the same units as fs. By
+            default, fs is 2 half-cycles/sample, so these are normalized
+            from 0 to 1, where 1 is the Nyquist frequency. (Wn is thus in
+            half-cycles / sample). For analog filters, Wn is an angular
+            frequency (e.g., rad/s). When Wn is a length-2 sequence,`Wn[0]`
+            must be less than `Wn[1]`.
+        rp:
+            For Chebyshev and elliptic filters, provides the maximum ripple in
+            the passband. (dB)
+        rs:
+            For Chebyshev and elliptic filters, provides the minimum
+            attenuation in the stop band. (dB)
+        btype:
+            The type of filter. Default is 'bandpass'.
+        analog:
+            When True, return an analog filter, otherwise a digital filter
+            is returned.
+        ftype:
+            The type of IIR filter to design:
+
+                - Butterworth   : 'butter'
+                - Chebyshev I   : 'cheby1'
+                - Chebyshev II  : 'cheby2'
+                - Cauer/elliptic: 'ellip'
+                - Bessel/Thomson: 'bessel's
+        fs:
+            The sampling frequency of the digital system.
+
+    Returns:
+        Filtered signal on the forward pass.
+    """
+
+    def __init__(
+        self,
+        N: int,
+        Wn: Union[float, torch.Tensor],
+        rs: Union[None, float, torch.Tensor] = None,
+        rp: Union[None, float, torch.Tensor] = None,
+        btype="band",
+        analog=False,
+        ftype="butter",
+        fs=None,
+    ) -> None:
+        super().__init__()
+
+        if isinstance(Wn, torch.Tensor):
+            Wn = Wn.numpy()
+        if isinstance(rs, torch.Tensor):
+            rs = rs.numpy()
+        if isinstance(rp, torch.Tensor):
+            rp = rp.numpy()
+
+        b, a = iirfilter(
+            N,
+            Wn,
+            rs=rs,
+            rp=rp,
+            btype=btype,
+            analog=analog,
+            ftype=ftype,
+            output="ba",
+            fs=fs,
+        )
+        self.register_buffer("b", torch.tensor(b))
+        self.register_buffer("a", torch.tensor(a))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        r"""
+        Apply the filter to the input signal.
+
+        Args:
+            x:
+                The input signal to be filtered.
+
+        Returns:
+            The filtered signal.
+        """
+        return filtfilt(x, self.a, self.b, clamp=False)

ml4gw/transforms/pearson.py

@@ -2,8 +2,8 @@ import torch
 from jaxtyping import Float
 from torch import Tensor
 
-from ml4gw.types import TimeSeries1to3d
-from ml4gw.utils.slicing import unfold_windows
+from ..types import TimeSeries1to3d
+from ..utils.slicing import unfold_windows
 
 
 class ShiftedPearsonCorrelation(torch.nn.Module):

ml4gw/transforms/qtransform.py

@@ -7,8 +7,8 @@ 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
+from ..types import FrequencySeries1to3d, TimeSeries1to3d, TimeSeries3d
+from .spline_interpolation import SplineInterpolate
 
 """
 All based on https://github.com/gwpy/gwpy/blob/v3.0.8/gwpy/signal/qtransform.py

ml4gw/transforms/scaler.py

@@ -4,7 +4,7 @@ import torch
 from jaxtyping import Float
 from torch import Tensor
 
-from ml4gw.transforms.transform import FittableTransform
+from .transform import FittableTransform
 
 
 class ChannelWiseScaler(FittableTransform):

ml4gw/transforms/snr_rescaler.py

@@ -2,9 +2,9 @@ from typing import Optional
 
 import torch
 
-from ml4gw.gw import compute_network_snr
-from ml4gw.transforms.transform import FittableSpectralTransform
-from ml4gw.types import BatchTensor, TimeSeries2d, WaveformTensor
+from ..gw import compute_network_snr
+from ..types import BatchTensor, TimeSeries2d, WaveformTensor
+from .transform import FittableSpectralTransform
 
 
 class SnrRescaler(FittableSpectralTransform):

ml4gw/transforms/spectral.py

@@ -4,8 +4,8 @@ import torch
 from jaxtyping import Float
 from torch import Tensor
 
-from ml4gw.spectral import fast_spectral_density, spectral_density
-from ml4gw.types import FrequencySeries1to3d, TimeSeries1to3d
+from ..spectral import fast_spectral_density, spectral_density
+from ..types import FrequencySeries1to3d, TimeSeries1to3d
 
 
 class SpectralDensity(torch.nn.Module):

ml4gw/transforms/spectrogram.py

@@ -7,7 +7,7 @@ from jaxtyping import Float
 from torch import Tensor
 from torchaudio.transforms import Spectrogram
 
-from ml4gw.types import TimeSeries3d
+from ..types import TimeSeries3d
 
 
 class MultiResolutionSpectrogram(torch.nn.Module):

ml4gw/transforms/transform.py

@@ -2,8 +2,8 @@ from typing import Optional
 
 import torch
 
-from ml4gw.spectral import spectral_density
-from ml4gw.types import FrequencySeries1to3d, TimeSeries1to3d
+from ..spectral import spectral_density
+from ..types import FrequencySeries1to3d, TimeSeries1to3d
 
 
 class FittableTransform(torch.nn.Module):

ml4gw/transforms/waveforms.py

@@ -4,8 +4,8 @@ import torch
 from jaxtyping import Float
 from torch import Tensor
 
-from ml4gw import gw
-from ml4gw.types import BatchTensor
+from .. import gw
+from ..types import BatchTensor
 
 
 # TODO: should these live in ml4gw.waveforms submodule?

ml4gw/transforms/whitening.py

@@ -2,14 +2,14 @@ from typing import Optional, Union
 
 import torch
 
-from ml4gw import spectral
-from ml4gw.transforms.transform import FittableSpectralTransform
-from ml4gw.types import (
+from .. import spectral
+from ..types import (
     FrequencySeries1d,
     FrequencySeries1to3d,
     TimeSeries1d,
     TimeSeries3d,
 )
+from .transform import FittableSpectralTransform
 
 
 class Whiten(torch.nn.Module):
@@ -45,6 +45,10 @@ class Whiten(torch.nn.Module):
             Cutoff frequency to apply highpass filtering
             during whitening. If left as `None`, no highpass
             filtering will be performed.
+        lowpass:
+            Cutoff frequency to apply lowpass filtering
+            during whitening. If left as `None`, no lowpass
+            filtering will be performed.
     """
 
     def __init__(
@@ -52,11 +56,13 @@ class Whiten(torch.nn.Module):
         fduration: float,
         sample_rate: float,
         highpass: Optional[float] = None,
+        lowpass: Optional[float] = None,
     ) -> None:
         super().__init__()
         self.fduration = fduration
         self.sample_rate = sample_rate
         self.highpass = highpass
+        self.lowpass = lowpass
 
         # register a window up front to signify our
         # fduration at inference time
@@ -104,6 +110,7 @@ class Whiten(torch.nn.Module):
             fduration=self.window,
             sample_rate=self.sample_rate,
             highpass=self.highpass,
+            lowpass=self.lowpass,
         )
 
 
@@ -153,6 +160,7 @@ class FixedWhiten(FittableSpectralTransform):
         *background: Union[TimeSeries1d, FrequencySeries1d],
         fftlength: Optional[float] = None,
         highpass: Optional[float] = None,
+        lowpass: Optional[float] = None,
         overlap: Optional[float] = None
     ) -> None:
         """
@@ -200,6 +208,13 @@ class FixedWhiten(FittableSpectralTransform):
                 in the frequency bins below this value to 0.
                 If left as `None`, the fit filter won't have any
                 highpass filtering properties.
+            lowpass:
+                Cutoff frequency, in Hz, used for lowpass filtering
+                with the fit whitening filter. This is achieved by
+                setting the frequency response of the fit PSDs
+                in the frequency bins above this value to 0.
+                If left as `None`, the fit filter won't have any
+                lowpass filtering properties.
             overlap:
                 Overlap between FFT frames used to convert
                 time-domain data to the frequency domain via
@@ -224,7 +239,7 @@ class FixedWhiten(FittableSpectralTransform):
             x = x.view(1, 1, -1)
 
             psd = spectral.truncate_inverse_power_spectrum(
-                x, fduration, self.sample_rate, highpass
+                x, fduration, self.sample_rate, highpass, lowpass
             )
             psds.append(psd[0, 0])
         psd = torch.stack(psds)

ml4gw/utils/slicing.py

@@ -5,12 +5,7 @@ from jaxtyping import Float, Int64
 from torch import Tensor
 from torch.nn.functional import unfold
 
-from ml4gw.types import (
-    TimeSeries1d,
-    TimeSeries1to3d,
-    TimeSeries2d,
-    TimeSeries3d,
-)
+from ..types import TimeSeries1d, TimeSeries1to3d, TimeSeries2d, TimeSeries3d
 
 BatchTimeSeriesTensor = Union[Float[Tensor, "batch time"], TimeSeries3d]
 

ml4gw/waveforms/cbc/coefficients.py

@@ -0,0 +1,35 @@
+import torch
+
+from ml4gw.constants import C, G
+from ml4gw.types import BatchTensor
+
+
+def taylor_t2_timing_0pn_coeff(total_mass: BatchTensor, eta: BatchTensor):
+    """
+    https://git.ligo.org/lscsoft/lalsuite/-/blob/master/lalsimulation/lib/LALSimInspiralPNCoefficients.c#L1528
+    """
+
+    output = total_mass * G / C**3
+    return -5.0 * output / (256.0 * eta)
+
+
+def taylor_t2_timing_2pn_coeff(eta: BatchTensor):
+    """
+    https://git.ligo.org/lscsoft/lalsuite/-/blob/master/lalsimulation/lib/LALSimInspiralPNCoefficients.c#L1545
+    """
+    return 7.43 / 2.52 + 11.0 / 3.0 * eta
+
+
+def taylor_t2_timing_4pn_coeff(eta: BatchTensor):
+    """
+    https://git.ligo.org/lscsoft/lalsuite/-/blob/master/lalsimulation/lib/LALSimInspiralPNCoefficients.c#L1560
+    """
+    return 30.58673 / 5.08032 + 54.29 / 5.04 * eta + 61.7 / 7.2 * eta**2
+
+
+def taylor_t3_frequency_0pn_coeff(total_mass: BatchTensor):
+    """
+    https://git.ligo.org/lscsoft/lalsuite/-/blob/master/lalsimulation/lib/LALSimInspiralPNCoefficients.c#L1723
+    """
+    output = total_mass * G / C**3.0
+    return 1.0 / (8.0 * torch.pi * output)

ml4gw/waveforms/cbc/phenom_d.py

@@ -1,9 +1,8 @@
 import torch
 from jaxtyping import Float
 
-from ml4gw.constants import MTSUN_SI, PI
-from ml4gw.types import BatchTensor, FrequencySeries1d
-
+from ...constants import MTSUN_SI, PI
+from ...types import BatchTensor, FrequencySeries1d
 from .phenom_d_data import QNMData_a, QNMData_fdamp, QNMData_fring
 from .taylorf2 import TaylorF2
 
@@ -26,6 +25,7 @@ class IMRPhenomD(TaylorF2):
         phic: BatchTensor,
         inclination: BatchTensor,
         f_ref: float,
+        **kwargs
     ):
         """
         IMRPhenomD waveform

ml4gw/waveforms/cbc/phenom_p.py

@@ -39,6 +39,7 @@ class IMRPhenomPv2(IMRPhenomD):
         inclination: BatchTensor,
         f_ref: float,
         tc: Optional[BatchTensor] = None,
+        **kwargs,
     ):
         """
         IMRPhenomPv2 waveform