ml4gw 0.7.5__py3-none-any.whl → 0.7.7__py3-none-any.whl

ml4gw/transforms/whitening.py

@@ -25,17 +25,17 @@ class Whiten(torch.nn.Module):
     In order to avoid edge effects due to filter settle-in,
     the provided PSDs will have their spectrum truncated
     such that their impulse response time in the time
-    domain is `fduration` seconds, and `fduration / 2`
+    domain is ``fduration`` seconds, and ``fduration / 2``
     seconds worth of data will be removed from each
     edge of the whitened timeseries.
 
-    For more information, see the documentation to
-    `ml4gw.spectral.whiten`.
+    For more information, see the documentation for
+    :meth:`~ml4gw.spectral.whiten`.
 
     Args:
         fduration:
             The length of the whitening filter's impulse
-            response, in seconds. `fduration / 2` seconds
+            response, in seconds. ``fduration / 2`` seconds
             worth of data will be cropped from the edges
             of the whitened timeseries.
         sample_rate:
@@ -43,11 +43,11 @@ class Whiten(torch.nn.Module):
             time is expected to be sampled
         highpass:
             Cutoff frequency to apply highpass filtering
-            during whitening. If left as `None`, no highpass
+            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
+            during whitening. If left as ``None``, no lowpass
             filtering will be performed.
     """
 
@@ -80,28 +80,28 @@ class Whiten(torch.nn.Module):
         Args:
             X:
                 Batch of multichannel timeseries to whiten.
-                Should have the shape (B, C, N), where
-                B is the batch size, C is the number of
-                channels, and N is the number of seconds
-                in the timeseries times `self.sample_rate`.
+                Should have the shape ``(B, C, N)``, where
+                ``B`` is the batch size, ``C`` is the number of
+                channels, and ``N`` is the number of seconds
+                in the timeseries times ``self.sample_rate``.
             psd:
                 Power spectral density used to whiten the
                 provided timeseries. Can be either 1D, 2D,
                 or 3D, with the last dimension representing
                 power at each frequency value. All other
                 dimensions must match their corresponding
-                value in `X`, starting from the right.
-                (e.g. if `psd.ndim == 2`, `psd.size(1)` should
-                be equal to `X.size(1)`. If `psd.ndim == 3`,
-                `psd.size(1)` and `psd.size(0)` should be equal
-                to `X.size(1)` and `X.size(0)`, respectively.)
+                value in ``X``, starting from the right.
+                (e.g. if ``psd.ndim == 2``, ``psd.size(1)`` should
+                be equal to ``X.size(1)``. If ``psd.ndim == 3``,
+                ``psd.size(1)`` and ``psd.size(0)`` should be equal
+                to ``X.size(1)`` and ``X.size(0)``, respectively.)
                 For more information about what these different
-                shapes for `psd` represent, consult the documentation
-                for `ml4gw.spectral.whiten`.
+                shapes for ``psd`` represent, consult the documentation
+                for :meth:`~ml4gw.spectral.whiten`.
         Returns:
-            Whitened timeseries, with `fduration * sample_rate / 2`
+            Whitened timeseries, with ``fduration * sample_rate / 2``
                 samples cropped from each edge. Output shape will then
-                be (B, C, N - `fduration * sample_rate`).
+                be ``(B, C, N - fduration * sample_rate)``.
         """
 
         return spectral.whiten(
@@ -118,7 +118,7 @@ class FixedWhiten(FittableSpectralTransform):
     """
     Transform that whitens timeseries by a fixed
     power spectral density that's determined by
-    calling the `.fit` method.
+    calling the ``.fit`` method.
 
     Args:
         num_channels:
@@ -167,7 +167,7 @@ class FixedWhiten(FittableSpectralTransform):
         Compute the PSD of channel-wise background to
         use to whiten timeseries at call time. PSDs will
         be resampled to have
-        `self.kernel_length * self.sample_rate // 2 + 1`
+        ``self.kernel_length * self.sample_rate // 2 + 1``
         frequency bins.
 
         Args:
@@ -176,29 +176,29 @@ class FixedWhiten(FittableSpectralTransform):
                 of the whitening filter, in seconds.
                 Fit PSDs will have their spectrum truncated
                 to approximate this response time.
-                A longer `fduration` will be able to
+                A longer ``fduration`` will be able to
                 handle narrower spikes in frequency, but
                 at the expense of longer filter settle-in
-                time. As such `fduration / 2` seconds of data
+                time. As such ``fduration / 2`` seconds of data
                 will be removed from each edge of whitened
                 timeseries.
             *background:
                 1D arrays capturing the signal to be used to
-                whiten each channel at call time. If `fftlength`
-                is left as `None`, it will be assumed that these
+                whiten each channel at call time. If ``fftlength``
+                is left as ``None``, it will be assumed that these
                 already represent frequency-domain data that will
                 be possibly resampled and truncated to whiten
                 timeseries at call time. Otherwise, it will be
                 assumed that these represent time-domain data that
                 will be converted to the frequency domain via
-                Welch's method using the specified `fftlength`
-                and `overlap`, with a Hann window used to window
+                Welch's method using the specified ``fftlength``
+                and ``overlap``, with a Hann window used to window
                 the FFT frames by default. Should have the same
-                number of args as `self.num_channels`.
+                number of args as ``self.num_channels``.
             fftlength:
                 Length of frames used to convert time-domain
                 data to the frequency-domain via Welch's method.
-                If left as `None`, it will be assumed that the
+                If left as ``None``, it will be assumed that the
                 background arrays passed already represent frequency-
                 domain data and don't require any conversion.
             highpass:
@@ -206,21 +206,21 @@ class FixedWhiten(FittableSpectralTransform):
                 with the fit whitening filter. This is achieved by
                 setting the frequency response of the fit PSDs
                 in the frequency bins below this value to 0.
-                If left as `None`, the fit filter won't have any
+                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
+                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
-                Welch's method. If `fftlength` is `None`, this
-                is ignored. Otherwise, if left as `None`, it will
-                be set to half of `fftlength` by default.
+                Welch's method. If ``fftlength`` is ``None``, this
+                is ignored. Otherwise, if left as ``None``, it will
+                be set to half of ``fftlength`` by default.
         """
         if len(background) != self.num_channels:
             raise ValueError(
@@ -250,8 +250,8 @@ class FixedWhiten(FittableSpectralTransform):
     def forward(self, X: TimeSeries3d) -> TimeSeries3d:
         """
         Whiten the input timeseries tensor using the
-        PSD fit by the `.fit` method, which must be
-        called _before_ the first call to `.forward`.
+        PSD fit by the ``.fit`` method, which must be
+        called **before** the first call to ``.forward``.
         """
         expected_dim = int(self.kernel_length * self.sample_rate)
         if X.size(-1) != expected_dim:

ml4gw/utils/slicing.py

@@ -25,9 +25,9 @@ def unfold_windows(
     Args:
         x:
             The timeseries to unfold. Can have shape
-            `(batch_size, num_channels, length * sample_rate)`,
-            `(num_channels, length * sample_rate)`, or
-            `(length * sample_rate)`
+            ``(batch_size, num_channels, length * sample_rate)``,
+            ``(num_channels, length * sample_rate)``, or
+            ``(length * sample_rate)``
         window_size:
             The size of the windows to unfold from x
         stride:
@@ -39,12 +39,12 @@ def unfold_windows(
 
     Returns:
        A tensor of shape
-       `(num_windows, batch_size, num_channels, kernel_size)`,
-       `(num_windows, num_channels, kernel_size)`, or
-       `(num_windows, kernel_size)` depending on whether the
+       ``(num_windows, batch_size, num_channels, kernel_size)``,
+       ``(num_windows, num_channels, kernel_size)``, or
+       ``(num_windows, kernel_size)`` depending on whether the
        input tensor is 3D, 2D, or 1D
 
-       If `drop_last` is false, returns the remainder of the
+       If ``drop_last`` is false, returns the remainder of the
        timeseries, shaped to be compatible with the returned
        unfolded tensor
     """
@@ -93,34 +93,34 @@ def slice_kernels(
     multichannel timeseries, slice kernels of a given size
     from the timeseries starting at the indicated indices.
     Returns a batch of 1D or 2D kernels, and so will have
-    one more dimension than `x`.
+    one more dimension than ``x``.
 
     Args:
         x:
             The timeseries tensor to slice kernels from
         idx:
-            The indices in `x` of the first sample of each
-            kernel. If `x` is 1D, `idx` must be 1D as well.
-            If `x` is 2D and `idx` is 1D, `idx` is assumed
+            The indices in ``x`` of the first sample of each
+            kernel. If ``x`` is 1D, ``idx`` must be 1D as well.
+            If ``x`` is 2D and ``idx`` is 1D, ``idx`` is assumed
             to represent the first index of the kernels sliced
             from _all_ channels (i.e. the channels are sliced
-            coincidentally). If `x` is 2D and `idx` is also 2D,
-            `idx` should have shape `(batch_size, num_channels)`,
+            coincidentally). If ``x`` is 2D and ``idx`` is also 2D,
+            ``idx`` should have shape ``(batch_size, num_channels)``,
             and its values are assumed to represent the first index
             of the kernel sliced from each channel _independently_.
-            If `x` is 3D, `idx` _must_ be 1D, and have the same length
-            as `x`. In this case, it is assumed that the elements of
-            `idx` represent the starting index in the last dimension
-            of `x` from which to sample a batch of kernels
+            If ``x`` is 3D, ``idx`` _must_ be 1D, and have the same length
+            as ``x``. In this case, it is assumed that the elements of
+            ``idx`` represent the starting index in the last dimension
+            of ``x`` from which to sample a batch of kernels
             coincidentally among the channels.
         kernel_size:
             The length of the kernels to slice from the timeseries
 
     Returns:
-        A tensor of shape `(batch_size, kernel_size)` if `x` is
-        1D and `(batch_size, num_channels, kernel_size)` if `x`
-        is 2D, where `batch_size = idx.shape[0]` and
-        `num_channels = x.shape[0]` if `x` is 2D.
+        A tensor of shape ``(batch_size, kernel_size)`` if ``x`` is
+        1D and ``(batch_size, num_channels, kernel_size)`` if ``x``
+        is 2D, where ``batch_size = idx.shape[0]`` and
+        ``num_channels = x.shape[0]`` if ``x`` is 2D.
     """
 
     # create the indices all the slices will be built around,
@@ -239,14 +239,14 @@ def sample_kernels(
 
     For a tensor representing one or multiple channels of
     timeseries data, randomly slice kernels of a fixed
-    length from the timeseries. If `X` is 1D, kernels will
-    be sampled uniformly from `X`. If `X` is 2D, kernels
-    will be sampled from the first dimension of `X` (assumed
+    length from the timeseries. If ``X`` is 1D, kernels will
+    be sampled uniformly from ``X``. If ``X`` is 2D, kernels
+    will be sampled from the first dimension of ``X`` (assumed
     to be the time dimension) in a manner that depends on the
-    values of the `max_center_offset` and `coincident` kwargs.
-    If `X` is 3D, one kernel will be sampled coincidentally from
-    each element along the 0th axis of `X`. In this case, `N` must
-    either be `None` or be equal to `len(X)`.
+    values of the ``max_center_offset`` and ``coincident`` kwargs.
+    If ``X`` is 3D, one kernel will be sampled coincidentally from
+    each element along the 0th axis of ``X``. In this case, ``N`` must
+    either be ``None`` or be equal to ``len(X)``.
 
     Args:
         X:
@@ -254,12 +254,12 @@ def sample_kernels(
             kernel_size: The size of the kernels to sample
         N:
             The number of kernels to sample. Can be left as
-            `None` if `X` is 3D, otherwise must be specified
+            ``None`` if ``X`` is 3D, otherwise must be specified
         max_center_offeset:
-            If `X` is 2D, this indicates the maximum distance
+            If ``X`` is 2D, this indicates the maximum distance
             from the center of the timeseries the edge of
-            sampled kernels may fall. If left as `None`, kernels
-            will be sampled uniformly across all of `X`'s time
+            sampled kernels may fall. If left as ``None``, kernels
+            will be sampled uniformly across all of ``X``'s time
             dimension. If greater than 0, defines the maximum
             distance that the rightmost edge of the kernel may
             fall from the center of the timeseries (the leftmost
@@ -269,19 +269,19 @@ def sample_kernels(
             of the timeseries, which may fall anywhere within the
             kernel with uniform probability. If less than 0, defines
             the minimum distance that the center of the timeseries
-            must fall from either edge of the kernel. If `X` is
+            must fall from either edge of the kernel. If ``X`` is
             1D, this argument is ignored.
         coincident:
-            If `X` is 2D, determines whether the individual channels
-            of `X` sample the same kernels or different kernels
+            If ``X`` is 2D, determines whether the individual channels
+            of ``X`` sample the same kernels or different kernels
             independently, i.e. whether the channels of each batch
             element in the output will contain coincident data. If
-            `X` is 1D, this argument is ignored.
+            ``X`` is 1D, this argument is ignored.
     Returns:
-        A batch of sampled kernels. If `X` is 1D, this will have
-        shape `(N, kernel_size)`. If `X` is 2D, this will have
-        shape `(N, num_channels, kernel_size)`, where
-        `num_channels = X.shape[0]`.
+        A batch of sampled kernels. If ``X`` is 1D, this will have
+        shape ``(N, kernel_size)``. If ``X`` is 2D, this will have
+        shape ``(N, num_channels, kernel_size)``, where
+        ``num_channels = X.shape[0]``.
     """
 
     if X.shape[-1] < kernel_size:

ml4gw/waveforms/cbc/phenom_d.py

@@ -56,14 +56,15 @@ class IMRPhenomD(TaylorF2):
         """
         # shape assumed (n_batch, params)
         if (
-            chirp_mass.shape[0] != mass_ratio.shape[0]
-            or mass_ratio.shape[0] != chi1.shape[0]
-            or chi1.shape[0] != chi2.shape[0]
-            or chi2.shape[0] != distance.shape[0]
-            or distance.shape[0] != phic.shape[0]
-            or phic.shape[0] != inclination.shape[0]
+            not chirp_mass.shape[0]
+            == mass_ratio.shape[0]
+            == chi1.shape[0]
+            == chi2.shape[0]
+            == distance.shape[0]
+            == phic.shape[0]
+            == inclination.shape[0]
         ):
-            raise RuntimeError("Tensors should have same batch size")
+            raise ValueError("Tensors must have same batch size")
         cfac = torch.cos(inclination)
         pfac = 0.5 * (1.0 + cfac * cfac)
 
@@ -104,16 +105,18 @@ class IMRPhenomD(TaylorF2):
 
         fRD, fDM = self.fring_fdamp(eta, eta2, chi1, chi2)
         Mf_peak = self.fmaxCalc(fRD, fDM, gamma2, gamma3)
-        _, t0 = self.phenom_d_mrd_phase(Mf_peak, eta, eta2, chi1, chi2, xi)
+        _, t0 = self.phenom_d_mrd_phase(
+            Mf_peak, eta, eta2, chi1, chi2, xi, fRD, fDM
+        )
 
         Mf = torch.outer(M_s, f)
         Mf_ref = torch.outer(M_s, f_ref * torch.ones_like(f))
 
         Psi, _ = self.phenom_d_phase(
-            Mf, mass_1, mass_2, eta, eta2, chi1, chi2, xi
+            Mf, mass_1, mass_2, eta, eta2, chi1, chi2, xi, fRD, fDM
         )
         Psi_ref, _ = self.phenom_d_phase(
-            Mf_ref, mass_1, mass_2, eta, eta2, chi1, chi2, xi
+            Mf_ref, mass_1, mass_2, eta, eta2, chi1, chi2, xi, fRD, fDM
         )
 
         Psi = (Psi.mT - 2 * phic).mT
@@ -133,6 +136,8 @@ class IMRPhenomD(TaylorF2):
             chi22,
             xi,
             distance,
+            fRD,
+            fDM,
         )
 
         amp_0 = self.taylorf2_amplitude(
@@ -157,8 +162,8 @@ class IMRPhenomD(TaylorF2):
         chi22,
         xi,
         distance,
-        fRD=None,  # used for passing ringdown frequency from phenom_p
-        fDM=None,  # used for passing damping frequency from phenom_p
+        fRD,
+        fDM,
     ):
         ins_amp, ins_Damp = self.phenom_d_inspiral_amp(
             Mf, eta, eta2, Seta, xi, chi1, chi2, chi12, chi22
@@ -173,14 +178,6 @@ class IMRPhenomD(TaylorF2):
         gamma2 = self.gamma2_fun(eta, eta2, xi)
         gamma3 = self.gamma3_fun(eta, eta2, xi)
 
-        # merger ringdown
-        if (fRD is None) != (fDM is None):
-            raise ValueError(
-                "Both fRD and fDM must either be provided or both be None"
-            )
-        if (fRD is None) and (fDM is None):
-            fRD, fDM = self.fring_fdamp(eta, eta2, chi1, chi2)
-
         Mf_peak = self.fmaxCalc(fRD, fDM, gamma2, gamma3)
         # Geometric peak and joining frequencies
         Mf_peak = (torch.ones_like(Mf).mT * Mf_peak).mT
@@ -221,17 +218,9 @@ class IMRPhenomD(TaylorF2):
         chi12,
         chi22,
         xi,
-        fRD=None,  # used for passing ringdown frequency from phenom_p
-        fDM=None,  # used for passing damping frequency from phenom_p
+        fRD,
+        fDM,
     ):
-        # merger ringdown
-        if (fRD is None) != (fDM is None):
-            raise ValueError(
-                "Both fRD and fDM must either be provided or both be None"
-            )
-        if (fRD is None) and (fDM is None):
-            fRD, fDM = self.fring_fdamp(eta, eta2, chi1, chi2)
-
         # Geometric frequency definition from PhenomD header file
         AMP_fJoin_INS = 0.014
 
@@ -268,19 +257,7 @@ class IMRPhenomD(TaylorF2):
         )
         return amp, Damp
 
-    # fRD and fDM are to be passed for generating phenom_p waveforms
-    # and remain None for phenom_d
-    def phenom_d_mrd_amp(
-        self, Mf, eta, eta2, chi1, chi2, xi, fRD=None, fDM=None
-    ):
-        # merger ringdown
-        if (fRD is None) != (fDM is None):
-            raise ValueError(
-                "Both fRD and fDM must either be provided or both be None"
-            )
-        if (fRD is None) and (fDM is None):
-            fRD, fDM = self.fring_fdamp(eta, eta2, chi1, chi2)
-
+    def phenom_d_mrd_amp(self, Mf, eta, eta2, chi1, chi2, xi, fRD, fDM):
         gamma1 = self.gamma1_fun(eta, eta2, xi)
         gamma2 = self.gamma2_fun(eta, eta2, xi)
         gamma3 = self.gamma3_fun(eta, eta2, xi)
@@ -422,10 +399,8 @@ class IMRPhenomD(TaylorF2):
 
         return amp, Damp
 
-    # fRD and fDM are to be passed for generating phenom_p waveforms
-    # and remain None for phenom_d
     def phenom_d_phase(
-        self, Mf, mass_1, mass_2, eta, eta2, chi1, chi2, xi, fRD=None, fDM=None
+        self, Mf, mass_1, mass_2, eta, eta2, chi1, chi2, xi, fRD, fDM
     ):
         ins_phase, ins_Dphase = self.phenom_d_inspiral_phase(
             Mf, mass_1, mass_2, eta, eta2, xi, chi1, chi2
@@ -435,13 +410,6 @@ class IMRPhenomD(TaylorF2):
             Mf, eta, eta2, chi1, chi2, xi, fRD, fDM
         )
 
-        # merger ringdown
-        if (fRD is None) != (fDM is None):
-            raise ValueError(
-                "Both fRD and fDM must either be provided or both be None"
-            )
-        if (fRD is None) and (fDM is None):
-            fRD, fDM = self.fring_fdamp(eta, eta2, chi1, chi2)
         # definitions in Eq. (35) of arXiv:1508.07253
         # PHI_fJoin_INS in header LALSimIMRPhenomD.h
         # C1 continuity at intermediate region i.e. f_1
@@ -501,25 +469,13 @@ class IMRPhenomD(TaylorF2):
 
         return phasing, Dphasing
 
-    # fRD and fDM are to be passed for generating phenom_p waveforms
-    # and remain None for phenom_d
-    def phenom_d_mrd_phase(
-        self, Mf, eta, eta2, chi1, chi2, xi, fRD=None, fDM=None
-    ):
+    def phenom_d_mrd_phase(self, Mf, eta, eta2, chi1, chi2, xi, fRD, fDM):
         alpha1 = self.alpha1Fit(eta, eta2, xi)
         alpha2 = self.alpha2Fit(eta, eta2, xi)
         alpha3 = self.alpha3Fit(eta, eta2, xi)
         alpha4 = self.alpha4Fit(eta, eta2, xi)
         alpha5 = self.alpha5Fit(eta, eta2, xi)
 
-        # merger ringdown
-        if (fRD is None) != (fDM is None):
-            raise ValueError(
-                "Both fRD and fDM must either be provided or both be None"
-            )
-        if (fRD is None) and (fDM is None):
-            fRD, fDM = self.fring_fdamp(eta, eta2, chi1, chi2)
-
         f_minus_alpha5_fRD = (Mf.t() - alpha5 * fRD).t()
         # Leading 1/eta is not multiplied at this stage
         mrd_phasing = (Mf.t() * alpha1).t()

ml4gw/waveforms/cbc/phenom_p.py

@@ -403,13 +403,17 @@ class IMRPhenomPv2(IMRPhenomD):
         with given data points :math:`(xp, fp)`, evaluated at :math:`x`
 
         Args:
-            x: the :math:`x`-coordinates at which to evaluate the interpolated
-                values.
-            xp: the :math:`x`-coordinates of data points, must be increasing.
-            fp: the :math:`y`-coordinates of data points, same length as `xp`.
+            x:
+                the :math:`x`-coordinates at which to evaluate the
+                interpolated values.
+            xp:
+                the :math:`x`-coordinates of data points, must be increasing.
+            fp:
+                the :math:`y`-coordinates of data points, same length as
+                ``xp``.
 
         Returns:
-            the interpolated values, same size as `x`.
+            the interpolated values, same size as ``x``.
         """
         original_shape = x.shape
         x = x.flatten()

ml4gw/waveforms/cbc/taylorf2.py

@@ -54,14 +54,15 @@ class TaylorF2(torch.nn.Module):
 
         # shape assumed (n_batch, params)
         if (
-            chirp_mass.shape[0] != mass_ratio.shape[0]
-            or chirp_mass.shape[0] != chi1.shape[0]
-            or chi1.shape[0] != chi2.shape[0]
-            or chi2.shape[0] != distance.shape[0]
-            or distance.shape[0] != phic.shape[0]
-            or phic.shape[0] != inclination.shape[0]
+            not chirp_mass.shape[0]
+            == mass_ratio.shape[0]
+            == chi1.shape[0]
+            == chi2.shape[0]
+            == distance.shape[0]
+            == phic.shape[0]
+            == inclination.shape[0]
         ):
-            raise RuntimeError("Tensors should have same batch size")
+            raise ValueError("Tensors must have same batch size")
 
         mass1, mass2 = chirp_mass_and_mass_ratio_to_components(
             chirp_mass, mass_ratio

ml4gw/waveforms/conversion.py

@@ -37,10 +37,11 @@ def chirp_mass_and_mass_ratio_to_components(
 ):
     """
     Compute component masses from chirp mass and mass ratio.
+
     Args:
         chirp_mass: Tensor of chirp mass values
         mass_ratio:
-            Tensor of mass ratio values, `m2 / m1`,
+            Tensor of mass ratio values, ``m2 / m1``,
             where m1 >= m2, so that mass_ratio <= 1
     """
     total_mass = chirp_mass * (1 + mass_ratio) ** 1.2 / mass_ratio**0.6