ml4gw 0.7.3__py3-none-any.whl → 0.7.4__py3-none-any.whl

ml4gw/dataloading/chunked_dataset.py

@@ -94,8 +94,10 @@ class ChunkedTimeSeriesDataset(torch.utils.data.IterableDataset):
             # flatten it to make it easier to slice
             if chunk_size < self.kernel_size:
                 raise ValueError(
-                    "Can't sample kernels of size {} from chunk "
-                    "with size {}".format(self.kernel_size, chunk_size)
+                    (
+                        "Can't sample kernels of size {} from chunk "
+                        "with size {}"
+                    ).format(self.kernel_size, chunk_size)
                 )
             chunk = chunk.reshape(-1)
 

ml4gw/dataloading/hdf5_dataset.py

@@ -105,6 +105,7 @@ class Hdf5TimeSeriesDataset(torch.utils.data.IterableDataset):
                             fname
                         ),
                         category=ContiguousHdf5Warning,
+                        stacklevel=2,
                     )
 
                 self.sizes[fname] = len(dset)

ml4gw/dataloading/in_memory_dataset.py

@@ -212,7 +212,6 @@ class InMemoryDataset(torch.utils.data.IterableDataset):
         Float[Tensor, "batch channel time"],
         Tuple[Float[Tensor, "batch channel time"], Float[Tensor, " batch"]],
     ]:
-
         indices = self.init_indices()
         for i in range(len(self)):
             # slice the array of _indices_ we'll be using to

ml4gw/distributions.py

@@ -4,6 +4,7 @@ from specified distributions. Each callable should map from
 an integer `N` to a 1D torch `Tensor` containing `N` samples
 from the corresponding distribution.
 """
+
 import math
 from typing import Optional
 
@@ -33,14 +34,15 @@ class Cosine(dist.Distribution):
         self.high = torch.as_tensor(high)
         self.norm = 1 / (torch.sin(self.high) - torch.sin(self.low))
 
-    def rsample(self, sample_shape: torch.Size = torch.Size()) -> Tensor:
+    def rsample(self, sample_shape: torch.Size = None) -> Tensor:
+        sample_shape = sample_shape or torch.Size()
         u = torch.rand(sample_shape, device=self.low.device)
         return torch.arcsin(u / self.norm + torch.sin(self.low))
 
     def log_prob(self, value: float) -> Float[Tensor, ""]:
         value = torch.as_tensor(value)
         inside_range = (value >= self.low) & (value <= self.high)
-        return value.cos().log() * inside_range
+        return (value.cos() * inside_range).log()
 
 
 class Sine(dist.TransformedDistribution):
@@ -166,7 +168,8 @@ class DeltaFunction(dist.Distribution):
         super().__init__(batch_shape, validate_args=validate_args)
         self.peak = torch.as_tensor(peak)
 
-    def rsample(self, sample_shape: torch.Size = torch.Size()) -> Tensor:
+    def rsample(self, sample_shape: torch.Size = None) -> Tensor:
+        sample_shape = sample_shape or torch.Size()
         return self.peak * torch.ones(
             sample_shape, device=self.peak.device, dtype=torch.float32
         )

ml4gw/gw.py

@@ -125,8 +125,8 @@ def compute_antenna_responses(
     for mode in modes:
         try:
             polarization = polarization_funcs[mode](m, n)
-        except KeyError:
-            raise ValueError(f"No polarization mode {mode}")
+        except KeyError as exc:
+            raise ValueError(f"No polarization mode {mode}") from exc
 
         # add a dummy dimension for concatenating
         polarizations.append(polarization)

ml4gw/nn/autoencoder/convolutional.py

@@ -19,7 +19,7 @@ class ConvBlock(Autoencoder):
         kernel_size: int,
         stride: int = 1,
         groups: int = 1,
-        activation: torch.nn.Module = torch.nn.ReLU(),
+        activation: torch.nn.Module = torch.nn.ReLU,
         norm: Module = torch.nn.BatchNorm1d,
         decode_channels: Optional[int] = None,
         output_activation: Optional[torch.nn.Module] = None,
@@ -56,7 +56,7 @@ class ConvBlock(Autoencoder):
             groups=groups,
         )
 
-        self.activation = activation
+        self.activation = activation()
         if output_activation is not None:
             self.output_activation = output_activation
         else:
@@ -97,7 +97,7 @@ class ConvolutionalAutoencoder(Autoencoder):
         kernel_size: int,
         stride: int = 1,
         groups: int = 1,
-        activation: torch.nn.Module = torch.nn.ReLU(),
+        activation: torch.nn.Module = torch.nn.ReLU,
         output_activation: Optional[torch.nn.Module] = None,
         norm: Module = torch.nn.BatchNorm1d,
         decode_channels: Optional[int] = None,

ml4gw/nn/resnet/resnet_1d.py

@@ -64,7 +64,6 @@ class BasicBlock(nn.Module):
         dilation: int = 1,
         norm_layer: Optional[Callable[..., nn.Module]] = None,
     ) -> None:
-
         super().__init__()
         if norm_layer is None:
             norm_layer = nn.BatchNorm1d
@@ -258,8 +257,10 @@ class ResNet1D(nn.Module):
             stride_type = ["stride"] * (len(layers) - 1)
         if len(stride_type) != (len(layers) - 1):
             raise ValueError(
-                "'stride_type' should be None or a "
-                "{}-element tuple, got {}".format(len(layers) - 1, stride_type)
+                (
+                    "'stride_type' should be None or a {}-element "
+                    "tuple, got {}"
+                ).format(len(layers) - 1, stride_type)
             )
 
         self.groups = groups

ml4gw/nn/resnet/resnet_2d.py

@@ -61,7 +61,6 @@ class BasicBlock(nn.Module):
         dilation: int = 1,
         norm_layer: Optional[Callable[..., nn.Module]] = None,
     ) -> None:
-
         super().__init__()
         if norm_layer is None:
             norm_layer = nn.BatchNorm2d
@@ -258,8 +257,10 @@ class ResNet2D(nn.Module):
             stride_type = ["stride"] * (len(layers) - 1)
         if len(stride_type) != (len(layers) - 1):
             raise ValueError(
-                "'stride_type' should be None or a "
-                "{}-element tuple, got {}".format(len(layers) - 1, stride_type)
+                (
+                    "'stride_type' should be None or a {}-element "
+                    "tuple, got {}"
+                ).format(len(layers) - 1, stride_type)
             )
 
         self.groups = groups

ml4gw/spectral.py

@@ -28,8 +28,8 @@ def median(x: Float[Tensor, "... size"], axis: int) -> Float[Tensor, "..."]:
     Implements a median calculation that matches numpy's
     behavior for an even number of elements and includes
     the same bias correction used by scipy's implementation.
-    see https://github.com/scipy/scipy/blob/main/scipy/signal/_spectral_py.py#L2066 # noqa
-    """
+    see https://github.com/scipy/scipy/blob/main/scipy/signal/_spectral_py.py#L2066
+    """  # noqa: E501
     n = x.shape[axis]
     ii_2 = 2 * torch.arange(1.0, (n - 1) // 2 + 1)
     bias = 1 + torch.sum(1.0 / (ii_2 + 1) - 1.0 / ii_2)
@@ -355,7 +355,7 @@ def truncate_inverse_power_spectrum(
     to which the whitening filter will be applied.
 
     Implementation details adapted from
-    https://github.com/vivinousi/gw-detection-deep-learning/blob/203966cc2ee47c32c292be000fb009a16824b7d9/modules/whiten.py#L8  # noqa
+    https://github.com/vivinousi/gw-detection-deep-learning/blob/203966cc2ee47c32c292be000fb009a16824b7d9/modules/whiten.py#L8
 
     Args:
         psd:
@@ -384,7 +384,7 @@ def truncate_inverse_power_spectrum(
         The PSD with its time domain response truncated
             to `fduration` and any filtered frequencies
             tapered.
-    """
+    """  # noqa: E501
 
     num_freqs = psd.size(-1)
     N = (num_freqs - 1) * 2
@@ -522,8 +522,10 @@ def whiten(
     N = X.size(-1)
     if N <= (2 * pad):
         raise ValueError(
-            "Not enough timeseries samples {} for "
-            "number of padded samples {}".format(N, 2 * pad)
+            (
+                "Not enough timeseries samples {} for number of "
+                "padded samples {}"
+            ).format(N, 2 * pad)
         )
 
     # normalize the number of expected dimensions in the PSD

ml4gw/transforms/iirfilter.py

@@ -10,7 +10,7 @@ class IIRFilter(torch.nn.Module):
     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
+    https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.iirfilter.html
 
     The forward call of this module accepts a batch tensor of shape
     (n_waveforms, n_samples) and returns the filtered waveforms.
@@ -50,7 +50,7 @@ class IIRFilter(torch.nn.Module):
 
     Returns:
         Filtered signal on the forward pass.
-    """
+    """  # noqa: E501
 
     def __init__(
         self,

ml4gw/transforms/qtransform.py

@@ -196,14 +196,14 @@ class SingleQTransform(torch.nn.Module):
         sample_rate: float,
         spectrogram_shape: Tuple[int, int],
         q: float = 12,
-        frange: List[float] = [0, torch.inf],
+        frange: List[float] = None,
         mismatch: float = 0.2,
         interpolation_method: str = "bicubic",
     ) -> None:
         super().__init__()
         self.q = q
         self.spectrogram_shape = spectrogram_shape
-        self.frange = frange
+        self.frange = frange or [0, torch.inf]
         self.duration = duration
         self.mismatch = mismatch
 
@@ -251,7 +251,7 @@ class SingleQTransform(torch.nn.Module):
         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)))
+        unique_ntiles = sorted(set(ntiles))
         idx = torch.arange(len(ntiles))
         self.stack_idx = [idx[Tensor(ntiles) == n] for n in unique_ntiles]
 
@@ -465,15 +465,15 @@ class QScan(torch.nn.Module):
         duration: float,
         sample_rate: float,
         spectrogram_shape: Tuple[int, int],
-        qrange: List[float] = [4, 64],
-        frange: List[float] = [0, torch.inf],
+        qrange: List[float] = None,
+        frange: List[float] = None,
         interpolation_method="bicubic",
         mismatch: float = 0.2,
     ) -> None:
         super().__init__()
-        self.qrange = qrange
+        self.qrange = qrange or [4, 64]
         self.mismatch = mismatch
-        self.frange = frange
+        self.frange = frange or [0, torch.inf]
         self.spectrogram_shape = spectrogram_shape
         max_q = torch.pi * duration * sample_rate / 50 - 11 ** (0.5)
         self.qs = self.get_qs()
@@ -481,7 +481,8 @@ class QScan(torch.nn.Module):
             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."
+                "scan will be truncated to avoid those values.",
+                stacklevel=2,
             )
 
         # Deliberately doing something different from GWpy here.

ml4gw/transforms/spectrogram.py

@@ -104,7 +104,8 @@ class MultiResolutionSpectrogram(torch.nn.Module):
         self.register_buffer("time_idxs", time_idxs)
 
     def _check_and_format_kwargs(self, kwargs: Dict[str, List]) -> List:
-        lengths = sorted(set([len(v) for v in kwargs.values()]))
+        lengths = sorted(len(v) for v in kwargs.values())
+        lengths = list(set(lengths))
 
         if lengths[-1] > 3:
             warnings.warn(
@@ -112,6 +113,7 @@ class MultiResolutionSpectrogram(torch.nn.Module):
                 "If performance is slower than desired, try reducing the "
                 "number of spectrograms",
                 RuntimeWarning,
+                stacklevel=2,
             )
 
         if len(lengths) > 2 or (len(lengths) == 2 and lengths[0] != 1):

ml4gw/transforms/spline_interpolation.py

@@ -74,7 +74,7 @@ class SplineInterpolate(torch.nn.Module):
     def __init__(
         self,
         x_in: Tensor,
-        y_in: Tensor = Tensor([1]),
+        y_in: Tensor = None,
         kx: int = 3,
         ky: int = 3,
         sx: float = 0.001,
@@ -83,6 +83,8 @@ class SplineInterpolate(torch.nn.Module):
         y_out: Optional[Tensor] = None,
     ):
         super().__init__()
+        if y_in is None:
+            y_in = Tensor([1])
         self.kx = kx
         self.ky = ky
         self.sx = sx
@@ -141,7 +143,6 @@ class SplineInterpolate(torch.nn.Module):
         d: int,
         m: int,
     ) -> Tuple[Tensor, Tensor]:
-
         """
         Compute the L and R values for B-spline basis functions.
         L and R are respectively the first and second coefficient multiplying
@@ -184,7 +185,6 @@ class SplineInterpolate(torch.nn.Module):
         n: int,
         m: int,
     ) -> Tensor:
-
         """
         Compute the zeroth-order B-spline basis functions
         according to de Boors recursive formula.
@@ -256,7 +256,6 @@ class SplineInterpolate(torch.nn.Module):
         return b[:, :, -1]
 
     def bivariate_spline_fit_natural(self, Z):
-
         if len(Z.shape) == 3:
             Z_Bx = torch.matmul(Z, self.Bx)
             # ((BxT @ Bx)^-1 @ (Z @ Bx)T)T = Z @ BxT^-1

ml4gw/transforms/whitening.py

@@ -161,7 +161,7 @@ class FixedWhiten(FittableSpectralTransform):
         fftlength: Optional[float] = None,
         highpass: Optional[float] = None,
         lowpass: Optional[float] = None,
-        overlap: Optional[float] = None
+        overlap: Optional[float] = None,
     ) -> None:
         """
         Compute the PSD of channel-wise background to

ml4gw/utils/slicing.py

@@ -171,8 +171,10 @@ def slice_kernels(
         # to select _different_ kernels from each channel
         if len(x) != idx.shape[1]:
             raise ValueError(
-                "Can't slice array with shape {} with indices "
-                "with shape {}".format(x.shape, idx.shape)
+                (
+                    "Can't slice array with shape {} with indices "
+                    "with shape {}"
+                ).format(x.shape, idx.shape)
             )
 
         # batch_size x num_channels x kernel_size
@@ -284,8 +286,9 @@ def sample_kernels(
 
     if X.shape[-1] < kernel_size:
         raise ValueError(
-            "Can't sample kernels of size {} from "
-            "tensor with shape {}".format(kernel_size, X.shape)
+            "Can't sample kernels of size {} from tensor with shape {}".format(
+                kernel_size, X.shape
+            )
         )
     elif X.ndim > 3:
         raise ValueError(
@@ -293,13 +296,14 @@ def sample_kernels(
         )
     elif X.ndim < 3 and N is None:
         raise ValueError(
-            "Must specify number of kernels N if X "
-            "has fewer than 3 dimensions"
+            "Must specify number of kernels N if X has fewer than 3 dimensions"
         )
     elif X.ndim == 3 and N is not None and N != len(X):
         raise ValueError(
-            "Can't sample {} kernels from 3D tensor "
-            "with batch dimension {}".format(N, len(X))
+            (
+                "Can't sample {} kernels from 3D tensor with "
+                "batch dimension {}"
+            ).format(N, len(X))
         )
 
     if X.ndim == 1:
@@ -309,7 +313,6 @@ def sample_kernels(
     center = int(X.shape[-1] // 2)
 
     if max_center_offset is None:
-
         # sample uniformly from all of X's time dimension
         min_val, max_val = 0, X.shape[-1] - kernel_size
     elif max_center_offset >= 0:
@@ -341,8 +344,10 @@ def sample_kernels(
         # if kernel_size > center - max_center_offset,
         # we may end up with negative indices
         raise ValueError(
-            "Kernel size {} is too large for requested center "
-            "offset value {}".format(kernel_size, max_center_offset)
+            (
+                "Kernel size {} is too large for requested center "
+                "offset value {}"
+            ).format(kernel_size, max_center_offset)
         )
 
     if X.ndim == 3 or coincident:

ml4gw/waveforms/cbc/phenom_d.py

@@ -25,7 +25,7 @@ class IMRPhenomD(TaylorF2):
         phic: BatchTensor,
         inclination: BatchTensor,
         f_ref: float,
-        **kwargs
+        **kwargs,
     ):
         """
         IMRPhenomD waveform

ml4gw/waveforms/cbc/phenom_p.py

@@ -1,6 +1,6 @@
 """
-    Based on the JAX implementation of IMRPhenomPv2 from
-    https://github.com/tedwards2412/ripple/blob/main/src/ripplegw/waveforms/IMRPhenomPv2.py
+Based on the JAX implementation of IMRPhenomPv2 from
+https://github.com/tedwards2412/ripple/blob/main/src/ripplegw/waveforms/IMRPhenomPv2.py
 """
 
 from typing import Dict, Optional, Tuple

ml4gw/waveforms/cbc/taylorf2.py

@@ -22,7 +22,7 @@ class TaylorF2(torch.nn.Module):
         phic: BatchTensor,
         inclination: BatchTensor,
         f_ref: float,
-        **kwargs
+        **kwargs,
     ):
         """
         TaylorF2 up to 3.5 PN in phase. Newtonian SPA amplitude.

ml4gw/waveforms/cbc/utils.py

@@ -1,7 +1,8 @@
 """
 Utilities for conditioning waveforms
-See https://git.ligo.org/lscsoft/lalsuite/-/blob/master/lalsimulation/lib/LALSimInspiral.c # noqa
-"""
+See https://git.ligo.org/lscsoft/lalsuite/-/blob/master/lalsimulation/lib/LALSimInspiral.c
+"""  # noqa: E501
+
 import torch
 
 from ml4gw.constants import MRSUN, MSUN, MTSUN_SI, C, G
@@ -105,7 +106,5 @@ def ringdown_time_bound(
 def frequency_isco(mass_1: BatchTensor, mass_2: BatchTensor):
     return (
         1.0
-        / (
-            (9.0**1.5) * torch.pi * (mass_1 + mass_2) * MTSUN_SI / MSUN
-        ).float()
+        / ((9.0**1.5) * torch.pi * (mass_1 + mass_2) * MTSUN_SI / MSUN).float()
     )

ml4gw/waveforms/conversion.py

@@ -20,15 +20,15 @@ def XLALSimInspiralLN(
     total_mass: BatchTensor, eta: BatchTensor, v: BatchTensor
 ):
     """
-    See https://git.ligo.org/lscsoft/lalsuite/-/blob/master/lalsimulation/lib/LALSimInspiralPNCoefficients.c#L2173 # noqa
-    """
+    See https://git.ligo.org/lscsoft/lalsuite/-/blob/master/lalsimulation/lib/LALSimInspiralPNCoefficients.c#L2173
+    """  # noqa: E501
     return total_mass**2 * eta / v
 
 
 def XLALSimInspiralL_2PN(eta: BatchTensor):
     """
-    See https://git.ligo.org/lscsoft/lalsuite/-/blob/master/lalsimulation/lib/LALSimInspiralPNCoefficients.c#L2181 # noqa
-    """
+    See https://git.ligo.org/lscsoft/lalsuite/-/blob/master/lalsimulation/lib/LALSimInspiralPNCoefficients.c#L2181
+    """  # noqa: E501
     return 1.5 + eta / 6.0
 
 
@@ -65,8 +65,8 @@ def bilby_spins_to_lalsim(
     """
     Converts between bilby spin and lalsimulation spin conventions.
 
-    See https://github.com/bilby-dev/bilby/blob/cccdf891e82d46319e69dbfdf48c4970b4e9a727/bilby/gw/conversion.py#L105 # noqa
-    and https://git.ligo.org/lscsoft/lalsuite/-/blob/master/lalsimulation/lib/LALSimInspiral.c#L3594 # noqa
+    See https://github.com/bilby-dev/bilby/blob/cccdf891e82d46319e69dbfdf48c4970b4e9a727/bilby/gw/conversion.py#L105
+    and https://git.ligo.org/lscsoft/lalsuite/-/blob/master/lalsimulation/lib/LALSimInspiral.c#L3594
 
     Args:
         theta_jn: BatchTensor,
@@ -80,13 +80,15 @@ def bilby_spins_to_lalsim(
         mass_2: BatchTensor,
         f_ref: float,
         phi_ref: BatchTensor,
-    """
+    """  # noqa: E501
 
     # check if f_ref is valid
     if f_ref <= 0.0:
         raise ValueError(
-            "f_ref <= 0 is invalid. "
-            "Please pass in the starting GW frequency instead."
+            (
+                "f_ref <= 0 is invalid. "
+                "Please pass in the starting GW frequency instead."
+            )
         )
 
     # starting frame: LNhat is along the z-axis and the unit

ml4gw/waveforms/generator.py

@@ -19,14 +19,15 @@ EXTRA_CYCLES = 3.0
 
 class TimeDomainCBCWaveformGenerator(torch.nn.Module):
     """
-    Waveform generator that generates time-domain waveforms from frequency-domain approximants.
+    Waveform generator that generates time-domain waveforms from
+    frequency-domain approximants.
 
     Frequency domain waveforms are conditioned as done by lalsimulation.
     Specifically, waveforms are generated with a starting frequency `fstart`
     slightly below the requested `f_min`, so that they can be tapered from
     `fstart` to `f_min` using a cosine window.
 
-    Please see https://lscsoft.docs.ligo.org/lalsuite/lalsimulation/group___l_a_l_sim_inspiral__c.html#gac9f16dab2cbca5a431738ee7d2505969 # noqa
+    Please see https://lscsoft.docs.ligo.org/lalsuite/lalsimulation/group___l_a_l_sim_inspiral__c.html#gac9f16dab2cbca5a431738ee7d2505969
     for more information
 
     Args:
@@ -50,7 +51,7 @@ class TimeDomainCBCWaveformGenerator(torch.nn.Module):
             will be placed.
         f_ref:
             Reference frequency for the waveform
-    """
+    """  # noqa: E501
 
     def __init__(
         self,
@@ -61,7 +62,6 @@ class TimeDomainCBCWaveformGenerator(torch.nn.Module):
         f_ref: float,
         right_pad: float,
     ) -> None:
-
         super().__init__()
         self.approximant = approximant
         self.f_min = f_min
@@ -93,8 +93,8 @@ class TimeDomainCBCWaveformGenerator(torch.nn.Module):
         """
         Builds highpass filter object.
 
-        See https://git.ligo.org/lscsoft/lalsuite/-/blob/master/lalsimulation/python/lalsimulation/gwsignal/core/conditioning_subroutines.py?ref_type=heads#L10 # noqa
-        """
+        See https://git.ligo.org/lscsoft/lalsuite/-/blob/master/lalsimulation/python/lalsimulation/gwsignal/core/conditioning_subroutines.py?ref_type=heads#L10
+        """  # noqa: E501
         order = 8.0
         w1 = np.tan(np.pi * (self.f_min) / self.sample_rate)
         attenuation = 0.99
@@ -117,22 +117,29 @@ class TimeDomainCBCWaveformGenerator(torch.nn.Module):
         self, **parameters: dict[str, BatchTensor]
     ) -> Tuple[Float[Tensor, "{N} samples"], Float[Tensor, "{N} samples"]]:
         """
-        Generate a conditioned frequency domain waveform from a frequency domain approximant.
+        Generate a conditioned frequency domain waveform from a
+        frequency-domain approximant.
 
-        Based on https://git.ligo.org/lscsoft/lalsuite/-/blob/master/lalsimulation/python/lalsimulation/gwsignal/core/waveform_conditioning.py?ref_type=heads#L248 # noqa
+        Based on https://git.ligo.org/lscsoft/lalsuite/-/blob/master/lalsimulation/python/lalsimulation/gwsignal/core/waveform_conditioning.py?ref_type=heads#L248
 
         Args:
             **parameters:
-                Dictionary of parameters for waveform generation
-                where key is the parameter name and value is a tensor of parameters.
-                It is required that `parameters` contains `mass_1`, `mass_2`, `s1z`, and `s2z`
-                keys, which are used for determining parameters of data conditioning.
-                If the specified approximant takes other parameters for waveform generation,
-                like `chirp_mass` and `mass_ratio`, the utility functions in `ml4gw.waveforms.conversion`
-                may be useful for populating the parameters dictionary with these additional parameters.
-                Note that, if using an approximant from `ml4gw.waveforms.cbc`, any additional keys in `parameters`
-                not ingested by the approximant will be ignored.
-        """
+                Dictionary of parameters for waveform generation where key is
+                the parameter name and value is a tensor of parameters.
+                It is required that `parameters` contains `mass_1`, `mass_2`,
+                `s1z`, and `s2z` keys, which are used for determining
+                parameters of data conditioning.
+
+                If the specified approximant takes other parameters for
+                waveform generation, like `chirp_mass` and `mass_ratio`, the
+                utility functions in `ml4gw.waveforms.conversion`may be useful
+                for populating the parameters dictionary with these
+                additional parameters.
+
+                Note that, if using an approximant from `ml4gw.waveforms.cbc`,
+                any additional keys in `parameters` not ingested by the
+                approximant will be ignored.
+        """  # noqa: E501
         # convert masses to kg, make sure
         # they are doubles so there is no
         # overflow in the calculations
@@ -266,23 +273,29 @@ class TimeDomainCBCWaveformGenerator(torch.nn.Module):
         **parameters,
     ) -> Tuple[Float[Tensor, "{N} samples"], Float[Tensor, "{N} samples"]]:
         """
-        Generates a time-domain waveform from a frequency domain approximant.
-        Conditioning is based onhttps://git.ligo.org/lscsoft/lalsuite/-/blob/master/lalsimulation/python/lalsimulation/gwsignal/core/waveform_conditioning.py?ref_type=heads#L248 # noqa
+        Generates a time-domain waveform from a frequency-domain approximant.
+        Conditioning is based onhttps://git.ligo.org/lscsoft/lalsuite/-/blob/master/lalsimulation/python/lalsimulation/gwsignal/core/waveform_conditioning.py?ref_type=heads#L248
 
-        A frequency domain waveform is generated, conditioned (see `generate_conditioned_fd_waveform`)
-        and fftdd into the time-domain
+        A frequency domain waveform is generated, conditioned
+        (see `generate_conditioned_fd_waveform`) and fft'd into the time-domain
 
         **parameters:
-            Dictionary of parameters for waveform generation
-            where key is the parameter name and value is a tensor of parameters.
-            It is required that `parameters` contains `mass_1`, `mass_2`, `s1z`, and `s2z`
-            keys, which are used for determining parameters of data conditioning.
-            If the specified approximant takes other parameters for waveform generation,
-            like `chirp_mass` and `mass_ratio`, the utility functions in `ml4gw.waveforms.conversion`
-            may be useful for populating the parameters dictionary with these additional parameters.
-            Note that, if using an approximant from `ml4gw.waveforms.cbc`, any additional keys in `parameters`
-            not ingested by the approximant will be ignored.
-        """
+            Dictionary of parameters for waveform generation where each key is
+            the parameter name and each value is a tensor of parameters.
+            It is required that `parameters` contains `mass_1`, `mass_2`,
+            `s1z`, and `s2z` keys, which are used for determining parameters
+            of data conditioning.
+
+            If the specified approximant takes other parameters for waveform
+            generation, like `chirp_mass` and `mass_ratio`, the utility
+            functions in `ml4gw.waveforms.conversion` may be useful for
+            populating the parameters dictionary with these additional
+            parameters.
+
+            Note that, if using an approximant from `ml4gw.waveforms.cbc`,
+            any additional keys in `parameters` not ingested by the
+            approximant will be ignored.
+        """  # noqa: E501
 
         hc, hp = self.generate_conditioned_fd_waveform(**parameters)
 

{ml4gw-0.7.3.dist-info → ml4gw-0.7.4.dist-info}/METADATA

@@ -1,20 +1,15 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.4
 Name: ml4gw
-Version: 0.7.3
+Version: 0.7.4
 Summary: Tools for training torch models on gravitational wave data
-Author: Alec Gunny
-Author-email: alec.gunny@ligo.org
-Requires-Python: >=3.9,<3.13
-Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.9
-Classifier: Programming Language :: Python :: 3.10
-Classifier: Programming Language :: Python :: 3.11
-Classifier: Programming Language :: Python :: 3.12
-Requires-Dist: jaxtyping (>=0.2,<0.3)
-Requires-Dist: numpy (<2.0.0)
-Requires-Dist: scipy (>=1.9.0,<1.15)
-Requires-Dist: torch (>=2.0,<3.0)
-Requires-Dist: torchaudio (>=2.0,<3.0)
+Author-email: Ethan Marx <emarx@mit.edu>, Will Benoit <benoi090@umn.edu>, Deep Chatterjee <deep1018@mit.edu>, Alec Gunny <alec.gunny@ligo.org>
+License-File: LICENSE
+Requires-Python: <3.13,>=3.9
+Requires-Dist: jaxtyping<0.3,>=0.2
+Requires-Dist: numpy<2.0.0
+Requires-Dist: scipy<1.15,>=1.9.0
+Requires-Dist: torchaudio~=2.0
+Requires-Dist: torch~=2.0
 Description-Content-Type: text/markdown
 
 # ML4GW
@@ -44,29 +39,6 @@ To build with a specific version of PyTorch/CUDA, please see the PyTorch install
 pip install ml4gw torch==2.5.1--extra-index-url=https://download.pytorch.org/whl/cu118
 ```
 
-### Poetry installation
-`ml4gw` is also fully compatible with use in Poetry, with your `pyproject.toml` set up like
-
-```toml
-[tool.poetry.dependencies]
-python = "^3.9"  # python versions 3.9-3.12 are supported
-ml4gw = "^0.6"
-```
-
-To build against a specific PyTorch/CUDA combination, consult the PyTorch installation documentation above and specify the `extra-index-url` via the `tool.poetry.source` table in your `pyproject.toml`. For example, to build against CUDA 11.6, you would do something like:
-
-```toml
-[tool.poetry.dependencies]
-python = "^3.9"
-ml4gw = "^0.6"
-torch = {version = "^2.0", source = "torch"}
-
-[[tool.poetry.source]]
-name = "torch"
-url = "https://download.pytorch.org/whl/cu118"
-priority = "explicit"
-```
-
 ## Contributing
 If you come across errors in the code, have difficulties using this software, or simply find that the current version doesn't cover your use case, please file an issue on our GitHub page, and we'll be happy to offer support.
 We encourage users who encounter these difficulties to file issues on GitHub, and we'll be happy to offer support to extend our coverage to new or improved functionality.
@@ -76,4 +48,3 @@ By bringing in new users with new use cases, we hope to develop this library int
 
 ## Funding
 We are grateful for the support of the U.S. National Science Foundation (NSF) Harnessing the Data Revolution (HDR) Institute for <a href="https://a3d3.ai">Accelerating AI Algorithms for Data Driven Discovery (A3D3)</a> under Cooperative Agreement No. <a href="https://www.nsf.gov/awardsearch/showAward?AWD_ID=2117997">PHY-2117997</a>.
-

{ml4gw-0.7.3.dist-info → ml4gw-0.7.4.dist-info}/RECORD

@@ -1,55 +1,55 @@
 ml4gw/__init__.py,sha256=81quoggCuIypZjZs3bbf1Ty70KHdva5RGEJxi0oC57E,25
 ml4gw/augmentations.py,sha256=pZH9tjEpXV0AIqvHHDkpUE-BorG02beOz2pmSipw2EY,1232
 ml4gw/constants.py,sha256=RQPXwavlw_cWu3ByltvTejPsi6EWXHDJQ1HaV9iE3Lg,850
+ml4gw/distributions.py,sha256=6UOgq8W-Bs-9170Jor_0hyeRnmC74zwbUrwAcJEz1jI,5082
+ml4gw/gw.py,sha256=0ovW_HJ3j2b5Yq3mduYtGLSl2RrvFyNNcOsZFf7koHY,19794
+ml4gw/spectral.py,sha256=sao_D0ceeMEatABfiabpqb-xxRfQO8Tz7yk9N7ciOAU,19858
+ml4gw/types.py,sha256=CcctqDcNajR7khGT6BD-WYsfRKpiP0udoSAB0k1qcFw,863
 ml4gw/dataloading/__init__.py,sha256=EHBBqU7y2-Np5iQ_xyufxamUEM1pPEquqFo7oaJnaJE,149
-ml4gw/dataloading/chunked_dataset.py,sha256=j96Rd67cRpsvotR_dzgfbrqxcoGDWnTV5cmfN038cb8,5256
-ml4gw/dataloading/hdf5_dataset.py,sha256=bVcXzS1LHVj7zMeMtRkxx1Q76MQS6wEApJJlUAI6iC8,7879
-ml4gw/dataloading/in_memory_dataset.py,sha256=1oUchfNBq3rx1NgNqrcg6AGdJ-dvm56o-TGFwPn5wm8,9546
-ml4gw/distributions.py,sha256=tUuaOiX5enjKLYWD7uiN8rdRVQcrIKps64xBkTl8fMs,4991
-ml4gw/gw.py,sha256=0I9MhoHWksWG9a5EUI0GkHD1skuOXiaQgSgxNKYXCxE,19778
+ml4gw/dataloading/chunked_dataset.py,sha256=yjJ_T7YFDErnLm-TmjXnZtT6GNPdPxj5njljWHr07Y0,5308
+ml4gw/dataloading/hdf5_dataset.py,sha256=IwOUYjjeUWPDL6HvPAfHrN5wK3qCrDdkH1OZFsecZnQ,7917
+ml4gw/dataloading/in_memory_dataset.py,sha256=7307LxdDGKxcy5_WNJp4f5J92AtV8gFT-B6P0QcWUfs,9545
 ml4gw/nn/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+ml4gw/nn/norm.py,sha256=JIOMXQbUtoWlrhncGsqW6f1-DiGDx9zQH2O3CvQml3U,3594
 ml4gw/nn/autoencoder/__init__.py,sha256=ZaT1XhJTHpMuPQqu5E__Jezeh9uwtjcXlT7IZ18byq4,161
 ml4gw/nn/autoencoder/base.py,sha256=eSWrDdpblI609oqa7RDSvZiY3YcV8WfhTioWKFn_7eE,3205
-ml4gw/nn/autoencoder/convolutional.py,sha256=VCemfNtzleIaKZHtzfxFDhYBLHdHEODy4-LRA9GiDhY,5359
+ml4gw/nn/autoencoder/convolutional.py,sha256=prj5Sat_yN15q-T7q9Uym8Hw6gjh7b8zimAVwT2Q85g,5357
 ml4gw/nn/autoencoder/skip_connection.py,sha256=9PKoCCvCUj5di9tuFM0Cl1v6gtcOK1bDeE_fS_R__FE,1391
 ml4gw/nn/autoencoder/utils.py,sha256=m_ivYGNwdrhA7cFxJVD4gqM8AHiWIGmlQI3pFNRklXQ,355
-ml4gw/nn/norm.py,sha256=JIOMXQbUtoWlrhncGsqW6f1-DiGDx9zQH2O3CvQml3U,3594
 ml4gw/nn/resnet/__init__.py,sha256=vBI0IftVP_EYAeDlqomtkGqUYE-RE_S4WNioUhniw9s,64
-ml4gw/nn/resnet/resnet_1d.py,sha256=C0H-GuY3-bradnSvUNtkD-o8j3-3uQDhUK4DKbOOrzk,13211
-ml4gw/nn/resnet/resnet_2d.py,sha256=fVzYRuO0xR9yGjjQExv30mouokvupOAW-Kfdbs5WYDA,13294
+ml4gw/nn/resnet/resnet_1d.py,sha256=qNiPNyT2Gracj0PvvRSwXiskqghZ22se5KYnlxcjzEw,13254
+ml4gw/nn/resnet/resnet_2d.py,sha256=xgEyraxgvAOM6OqEtpciRCa20U_fyClLG-A8oaIr2UQ,13337
 ml4gw/nn/streaming/__init__.py,sha256=zgjGR2L8t0txXLnil9ceZT0tM8Y2FC8yPxqIKYH0o1A,80
 ml4gw/nn/streaming/online_average.py,sha256=_nrul4ygTC_ln4wpSWGRWTgWlfGeOUGXxeGrhU4oJms,4716
 ml4gw/nn/streaming/snapshotter.py,sha256=1vWDpebRQBZIUVeksbXoqngqMnlSzQFkcsgYNrHB9tc,4473
-ml4gw/spectral.py,sha256=rnxd1ObPjyQMAu3D83_sw2lEEHZF7f87YQBV_pxHLxM,19809
 ml4gw/transforms/__init__.py,sha256=OaTQJD4GFkDkcxt0DIwt2AzeEcv9t21ciKXxQnqDiuI,447
-ml4gw/transforms/iirfilter.py,sha256=RwgC3DWgYmBnHe7bYjvr9njM1WrRZ9EjBJsZNmtOY8s,3186
+ml4gw/transforms/iirfilter.py,sha256=0FQ2HfXiYuvRzNmMz-SxkocTTrrKuJpxob0NGgkalBA,3188
 ml4gw/transforms/pearson.py,sha256=CM9FTRxI4384-36FIaJFOcMZwsA7BkgberToJkMU1PA,3227
-ml4gw/transforms/qtransform.py,sha256=5S9y3PxkOmqMAarQmme0Tiy58vRvberpqhg6IeyDJLI,20675
+ml4gw/transforms/qtransform.py,sha256=Kntl3MlZvL5Ae4A4Vskfq-uIHKhue_USnEHLG9Cf-Pw,20724
 ml4gw/transforms/scaler.py,sha256=K5mp4w2zGZbpH1AcBUfpQS4n3aVSNzkaGWXedwk2LXs,2508
 ml4gw/transforms/snr_rescaler.py,sha256=lfuwdwMY117gB-emmn0_22gsK_A9xnkHJv2-76HFWc4,2728
 ml4gw/transforms/spectral.py,sha256=4uCLNEcDff4kLheUA5v64L0y_MSOvUTJ92IH4TVcEys,4385
-ml4gw/transforms/spectrogram.py,sha256=8HDStoup7vlwpw9qTKshAuEpa85-lw5_SwYxjxxu1sQ,6158
-ml4gw/transforms/spline_interpolation.py,sha256=oXih-gLMbIbI36DPKLTk39WcjiNUJN_rcQia_k3OrMY,13527
+ml4gw/transforms/spectrogram.py,sha256=fpI0vypB5seBA05a979KKsSL9go0FHGI_-BVzRJ9ATA,6218
+ml4gw/transforms/spline_interpolation.py,sha256=CDfHX2Cd0KqNgo2JXNk2QgFNYtDfEXCMIFZ7bDqo6po,13573
 ml4gw/transforms/transform.py,sha256=lu5ukcOCOYYZDZCM_0amS9AY2bJgkbLpXmZ9DpnSK9I,2504
 ml4gw/transforms/waveforms.py,sha256=koWOuHuUpQWmTT1yawSWa_MOuLfDBuugy91KIyuklOo,3189
-ml4gw/transforms/whitening.py,sha256=8ADmM52lrHt_2yvjX51x0bFxAloKbS7s2owJgrVD5uc,10294
-ml4gw/types.py,sha256=CcctqDcNajR7khGT6BD-WYsfRKpiP0udoSAB0k1qcFw,863
+ml4gw/transforms/whitening.py,sha256=XyBAAsTkW5O4K06ET6a9mb0809yDkMcVZJjkB1X1vY4,10295
 ml4gw/utils/interferometer.py,sha256=lRS0N3SwUTknhYXX57VACJ99jK1P9M19oUWN_i_nQN0,1814
-ml4gw/utils/slicing.py,sha256=V9tbEzHnukg16-e8jdIFsZIZ1oICF9zBE2sjUsBXW-s,13538
+ml4gw/utils/slicing.py,sha256=WkY9biz6V8FWUmqCq9QpYZv67kp0yEjdKfkI8CcK4Dc,13653
 ml4gw/waveforms/__init__.py,sha256=QVUzBx_y8A9_AsRuTJruPvL9mqGnBt11Iw1MOYjXyE4,40
+ml4gw/waveforms/conversion.py,sha256=YMyjxzPlLAay66oNABCU98VnvhyjTNbGJ36R5sX3MVo,6945
+ml4gw/waveforms/generator.py,sha256=RuVO9FL3IwmzTb10JkQ4iR8u2-0edy8DiZhZ_leHAPc,12169
 ml4gw/waveforms/adhoc/__init__.py,sha256=XVwP4t8TMUj87WY3yMGRTkXsv7_lVr1w8p8iKBW8iKE,71
 ml4gw/waveforms/adhoc/ringdown.py,sha256=m8IBQTxKBBGFqBtWGEO4KG3DEYR8TTnNyGVdVLaMKa8,3316
 ml4gw/waveforms/adhoc/sine_gaussian.py,sha256=-MtrI7ydwBTk4K0O4tdkC8-w5OifQszdnWN9__I4XzY,3569
 ml4gw/waveforms/cbc/__init__.py,sha256=hGbPsFNAIveYJnff8qKY8RWeBPFtZoYcnGHxraPWtWI,99
 ml4gw/waveforms/cbc/coefficients.py,sha256=PMr0IBALEQ38eAvZqYg-w_FE_sS1mH2FWr9soQ5MRfU,1106
-ml4gw/waveforms/cbc/phenom_d.py,sha256=b586PbpBGAA1DO55X0D35_dAJXIGVwUBrNhmPgCBbwU,48983
+ml4gw/waveforms/cbc/phenom_d.py,sha256=GdBdSfUIze2KUfKBS44A4kfYpG27aA5FYknaBqL0q8c,48984
 ml4gw/waveforms/cbc/phenom_d_data.py,sha256=WA1FBxUp9fo1IQaV_OLJ_5g5gI166mY1FtG9n25he9U,53447
-ml4gw/waveforms/cbc/phenom_p.py,sha256=LBtGVUjBjROcYBPLldFnF6T1jZV6ZyuZEnkn9-oTKpQ,27620
-ml4gw/waveforms/cbc/taylorf2.py,sha256=2ga_lG_xkYOsF-BdxgjbU0pgLDjeAO0p5IWuCPvibvQ,10504
-ml4gw/waveforms/cbc/utils.py,sha256=CvZ79PQygb-zwulMV-wRuBcGEsHbVOtJz60UnOJFKoM,3051
-ml4gw/waveforms/conversion.py,sha256=MyADWEZVoEkRkKaHk1ZuQDsGfPYx5xUTtyApj5P3ueQ,6895
-ml4gw/waveforms/generator.py,sha256=i2lgaJzH5eA6gzc-bLQZYYEgEQ8OBLJgE9yNXU3FsKM,12005
-ml4gw-0.7.3.dist-info/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
-ml4gw-0.7.3.dist-info/METADATA,sha256=LSSAqGyEFUoINHjSh1iCc3TRRY-ytTU6pPZ6TVeZdLI,3941
-ml4gw-0.7.3.dist-info/WHEEL,sha256=IYZQI976HJqqOpQU6PHkJ8fb3tMNBFjg-Cn-pwAbaFM,88
-ml4gw-0.7.3.dist-info/RECORD,,
+ml4gw/waveforms/cbc/phenom_p.py,sha256=RZzzKQzqZW3rQuWZ41htTZOwwulYP61ow87HRRrel5A,27612
+ml4gw/waveforms/cbc/taylorf2.py,sha256=cmYrVL29dwX2Icp7I6SXqRIjtPmoljK5DP_ofx2heiM,10505
+ml4gw/waveforms/cbc/utils.py,sha256=LT1ky10_6ZrbwTcxIrWP1O75GUEuU5q2ZE2yYDhadQE,3037
+ml4gw-0.7.4.dist-info/METADATA,sha256=5nM8sBFDpqrKHQNqtssbhFYFkwze3IE_HLM8Zb8qXQU,3049
+ml4gw-0.7.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+ml4gw-0.7.4.dist-info/licenses/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
+ml4gw-0.7.4.dist-info/RECORD,,

{ml4gw-0.7.3.dist-info → ml4gw-0.7.4.dist-info}/WHEEL

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