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

ml4gw/gw.py

@@ -2,13 +2,11 @@
 Tools for manipulating raw gravitational waveforms
 and projecting them onto interferometer responses.
 Much of the projection code is an extension of the
-implementation made available in bilby:
-
-https://arxiv.org/abs/1811.02042
-
-Specifically the code here:
-https://github.com/lscsoft/bilby/blob/master/bilby/gw/detector/interferometer.py
-"""
+implementation made available in
+`bilby <https://arxiv.org/abs/1811.02042>`_.
+Specifically code from
+`this module <https://github.com/lscsoft/bilby/blob/master/bilby/gw/detector/interferometer.py>`_.
+"""  # noqa E501
 
 from typing import List, Tuple, Union
 
@@ -134,6 +132,9 @@ def compute_antenna_responses(
     # shape: batch x num_polarizations x 3 x 3
     polarization = torch.stack(polarizations, axis=1)
 
+    # Ensure dtype consistency before einsum
+    detector_tensors = detector_tensors.to(polarization.dtype)
+
     # compute the weight of each interferometer's response
     # to each polarization: batch x polarizations x ifos
     return torch.einsum("...jk,ijk->...i", polarization, detector_tensors)
@@ -194,7 +195,7 @@ def compute_observed_strain(
     **polarizations: Float[Tensor, "batch time"],
 ) -> WaveformTensor:
     """
-    Compute the strain timeseries $h(t)$ observed by a network
+    Compute the strain timeseries :math:`h(t)` observed by a network
     of interferometers from the given polarization timeseries
     corresponding to gravitational waveforms from sources with
     the indicated sky parameters.
@@ -222,13 +223,13 @@ def compute_observed_strain(
             between the waveform observed at the geocenter and
             the one observed at the detector site. To avoid
             adding any delay between the two, reset your coordinates
-            such that the desired interferometer is at `(0., 0., 0.)`.
+            such that the desired interferometer is at ``(0., 0., 0.)``.
         sample_rate:
             Rate at which the polarization timeseries have been sampled
         polarziations:
             Timeseries for each waveform polarization which
             contributes to the interferometer response. Allowed
-            polarizations are `cross`, `plus`, and `breathing`.
+            polarizations are ``cross``, ``plus``, and ``breathing``.
     Returns:
         Tensor representing the observed strain at each
         interferometer for each waveform.
@@ -236,13 +237,15 @@ def compute_observed_strain(
 
     # TODO: just use theta as the input parameter?
     # note that ** syntax is ordered, so we're safe
-    # to be lazy and use `list` for the keys and values
+    # to be lazy and use ``list`` for the keys and values
     theta = torch.pi / 2 - dec
     antenna_responses = compute_antenna_responses(
         theta, psi, phi, detector_tensors, list(polarizations)
     )
 
     polarizations = torch.stack(list(polarizations.values()), axis=1)
+    # Ensure dtype consistency before einsum
+    antenna_responses = antenna_responses.to(polarizations.dtype)
     waveforms = torch.einsum(
         "...pi,...pt->...it", antenna_responses, polarizations
     )
@@ -286,26 +289,28 @@ def compute_ifo_snr(
     highpass: Union[float, Float[Tensor, " frequency"], None] = None,
     lowpass: Union[float, Float[Tensor, " frequency"], None] = None,
 ) -> Float[Tensor, "batch num_ifos"]:
-    r"""Compute the SNRs of a batch of interferometer responses
+    """Compute the SNRs of a batch of interferometer responses
 
     Compute the signal to noise ratio (SNR) of individual
     interferometer responses to gravitational waveforms with
     respect to a background PSD for each interferometer. The
-    SNR of the $i$th waveform at the $j$th interferometer
+    SNR of the :math:`i` th waveform at the :math:`j` th interferometer
     is computed as:
 
-    $$\rho_{ij} =
-        4 \int_{f_{\text{min}}}^{f_{\text{max}}}
-        \frac{\tilde{h_{ij}}(f)\tilde{h_{ij}}^*(f)}
-        {S_n^{(j)}(f)}df$$
+    .. math::
 
-    Where $f_{\text{min}}$ is a minimum frequency denoted
-    by `highpass`, `f_{\text{max}}` is the maximum frequency
-    denoted by `lowpass`, which defaults to the Nyquist frequency
-    dictated by `sample_rate`; `\tilde{h_{ij}}` and `\tilde{h_{ij}}*`
-    indicate the fourier transform of the $i$th waveform at
-    the $j$th inteferometer and its complex conjugate, respectively;
-    and $S_n^{(j)}$ is the backround PSD at the $j$th interferometer.
+        \\rho_{ij} =
+        4 \\int_{f_{\\text{min}}}^{f_{\\text{max}}}
+        \\frac{\\tilde{h_{ij}}(f)\\tilde{h_{ij}}^*(f)}
+        {S_n^{(j)}(f)}df
+
+    Where :math:`f_{\\text{min}}` is a minimum frequency denoted
+    by ``highpass``, :math:`f_{\\text{max}}` is the maximum frequency
+    denoted by ``lowpass``, which defaults to the Nyquist frequency
+    dictated by ``sample_rate``; :math:`\\tilde{h}_{ij}` and :math:`\\tilde{h}_{ij}^*`
+    indicate the fourier transform of the :math:`i` th waveform at
+    the :math:`j` th inteferometer and its complex conjugate, respectively;
+    and :math:`S_n^{(j)}` is the backround PSD at the :math:`j` th interferometer.
 
     Args:
         responses:
@@ -314,12 +319,12 @@ def compute_ifo_snr(
         psd:
             The one-sided power spectral density of the background
             noise at each interferometer to which a response
-            in `responses` has been calculated. If 2D, each row of
-            `psd` will be assumed to be the background PSD for each
-            channel of _every_ batch element in `responses`. If 3D,
+            in ``responses`` has been calculated. If 2D, each row of
+            ``psd`` will be assumed to be the background PSD for each
+            channel of _every_ batch element in ``responses``. If 3D,
             this should contain a background PSD for each channel
-            of each element in `responses`, and therefore the first
-            two dimensions of `psd` and `responses` should match.
+            of each element in ``responses``, and therefore the first
+            two dimensions of ``psd`` and ``responses`` should match.
         sample_rate:
             The frequency at which the waveform responses timeseries
             have been sampled. Upon fourier transforming, should
@@ -329,18 +334,18 @@ def compute_ifo_snr(
             If a tensor is provided, it will be assumed to be a
             pre-computed mask used to 0-out low frequency components.
             If a float, it will be used to compute such a mask. If
-            left as `None`, all frequencies up to `lowpass`
+            left as ``None``, all frequencies up to ``lowpass``
             will contribute to the SNR calculation.
         lowpass:
             The maximum frequency below which to compute the SNR.
             If a tensor is provided, it will be assumed to be a
             pre-computed mask used to 0-out high frequency components.
             If a float, it will be used to compute such a mask. If
-            left as `None`, all frequencies from `highpass` up to
+            left as ``None``, all frequencies from ``highpass`` up to
             the Nyquist freqyency will contribute to the SNR calculation.
     Returns:
         Batch of SNRs computed for each interferometer
-    """
+    """  # noqa E501
 
     # TODO: should we do windowing here?
     # compute frequency power, upsampling precision so that
@@ -388,10 +393,10 @@ def compute_ifo_snr(
     # that the user specify the sample rate by taking the
     # fft as-is (without dividing by sample rate) and then
     # taking the mean here (or taking the sum and dividing
-    # by the sum of `highpass` if it's a mask). If we want
+    # by the sum of ``highpass`` if it's a mask). If we want
     # to allow the user to pass a float for highpass, we'll
     # need the sample rate to compute the mask, but if we
-    # replace this with a `mask` argument instead we're in
+    # replace this with a ``mask`` argument instead we're in
     # the clear
     df = sample_rate / responses.shape[-1]
     integrated = integrand.sum(axis=-1) * df
@@ -408,15 +413,17 @@ def compute_network_snr(
     highpass: Union[float, Float[Tensor, " frequency"], None] = None,
     lowpass: Union[float, Float[Tensor, " frequency"], None] = None,
 ) -> BatchTensor:
-    r"""
+    """
     Compute the total SNR from a gravitational waveform
     from a network of interferometers. The total SNR for
-    the $i$th waveform is computed as
+    the :math:`i` th waveform is computed as
+
+    .. math::
 
-    $$\rho_i = \sqrt{\sum_{j}^{N}\rho_{ij}^2}$$
+        \\rho_i = \\sqrt{\\sum_{j}^{N}\\rho_{ij}^2}
 
-    where \rho_{ij} is the SNR for the $i$th waveform at
-    the $j$th interferometer in the network and $N$ is
+    where :math:`\\rho_{ij}` is the SNR for the :math:`i` th waveform at
+    the :math:`j` th interferometer in the network and :math:`N` is
     the total number of interferometers.
 
     Args:
@@ -426,12 +433,12 @@ def compute_network_snr(
         backgrounds:
             The one-sided power spectral density of the background
             noise at each interferometer to which a response
-            in `responses` has been calculated. If 2D, each row of
-            `psd` will be assumed to be the background PSD for each
-            channel of _every_ batch element in `responses`. If 3D,
+            in ``responses`` has been calculated. If 2D, each row of
+            ``psd`` will be assumed to be the background PSD for each
+            channel of **every** batch element in ``responses``. If 3D,
             this should contain a background PSD for each channel
-            of each element in `responses`, and therefore the first
-            two dimensions of `psd` and `responses` should match.
+            of each element in ``responses``, and therefore the first
+            two dimensions of ``psd`` and ``responses`` should match.
         sample_rate:
             The frequency at which the waveform responses timeseries
             have been sampled. Upon fourier transforming, should
@@ -441,14 +448,14 @@ def compute_network_snr(
             If a tensor is provided, it will be assumed to be a
             pre-computed mask used to 0-out low frequency components.
             If a float, it will be used to compute such a mask. If
-            left as `None`, all frequencies up to `sample_rate / 2`
+            left as ``None``, all frequencies up to ``sample_rate / 2``
             will contribute to the SNR calculation.
         lowpass:
             The maximum frequency below which to compute the SNR.
             If a tensor is provided, it will be assumed to be a
             pre-computed mask used to 0-out high frequency components.
             If a float, it will be used to compute such a mask. If
-            left as `None`, all frequencies from `highpass` up to
+            left as ``None``, all frequencies from ``highpass`` up to
             the Nyquist freqyency will contribute to the SNR calculation.
     Returns:
         Batch of SNRs for each waveform across the interferometer network
@@ -478,12 +485,12 @@ def reweight_snrs(
         psd:
             The one-sided power spectral density of the background
             noise at each interferometer to which a response
-            in `responses` has been calculated. If 2D, each row of
-            `psd` will be assumed to be the background PSD for each
-            channel of _every_ batch element in `responses`. If 3D,
+            in ``responses`` has been calculated. If 2D, each row of
+            ``psd`` will be assumed to be the background PSD for each
+            channel of **every** batch element in ``responses``. If 3D,
             this should contain a background PSD for each channel
-            of each element in `responses`, and therefore the first
-            two dimensions of `psd` and `responses` should match.
+            of each element in ``responses``, and therefore the first
+            two dimensions of ``psd`` and ``responses`` should match.
         sample_rate:
             The frequency at which the waveform responses timeseries
             have been sampled. Upon fourier transforming, should
@@ -493,14 +500,14 @@ def reweight_snrs(
             If a tensor is provided, it will be assumed to be a
             pre-computed mask used to 0-out low frequency components.
             If a float, it will be used to compute such a mask. If
-            left as `None`, all frequencies up to `sample_rate / 2`
+            left as ``None``, all frequencies up to ``sample_rate / 2``
             will contribute to the SNR calculation.
         lowpass:
             The maximum frequency below which to compute the SNR.
             If a tensor is provided, it will be assumed to be a
             pre-computed mask used to 0-out high frequency components.
             If a float, it will be used to compute such a mask. If
-            left as `None`, all frequencies from `highpass` up to
+            left as ``None``, all frequencies from ``highpass`` up to
             the Nyquist freqyency will contribute to the SNR calculation.
     Returns:
         Rescaled interferometer responses

ml4gw/nn/autoencoder/base.py

@@ -12,18 +12,18 @@ class Autoencoder(torch.nn.Module):
     Base autoencoder class that defines some of the
     basic methods and functionality. Autoencoders are
     defined here as a set of sequential blocks that
-    have an `encode` method, which acts on the input
-    data to the autoencoder, and a `decode` method, which
-    acts on the encoded vector generated by the `encode`
-    method. `forward` just runs these steps one after the
+    have an ``encode`` method, which acts on the input
+    data to the autoencoder, and a ``decode`` method, which
+    acts on the encoded vector generated by the ``encode``
+    method. ``forward`` just runs these steps one after the
     other. Although it isn't explicitly enforced, a good
-    rule of thumb is that the ouput of a block's `decode`
+    rule of thumb is that the ouput of a block's ``decode``
     method should have the same shape as the _input_ of its
-    `encode` method.
+    ``encode`` method.
 
-    Accepts a `skip_connection` argument that defines how to
-    combine information from the input of one block's `encode`
-    layer with the output to its `decode`layer. See `skip_connections.py`
+    Accepts a ``skip_connection`` argument that defines how to
+    combine information from the input of one block's ``encode``
+    layer with the output to its ``decode`` layer. See ``skip_connections.py``
     for more info about what these classes are expected to contain
     and how they operate.
     """

ml4gw/nn/autoencoder/convolutional.py

@@ -83,11 +83,11 @@ class ConvolutionalAutoencoder(Autoencoder):
     match the shape of the input to its corresponding
     encoder layer, except for the last decoder which
     can have an arbitrary number of channels specified
-    by `decode_channels`.
+    by ``decode_channels``.
 
-    All layers also share the same `activation` except
+    All layers also share the same ``activation`` except
     for the last decoder layer, which can have an
-    arbitrary `output_activation`.
+    arbitrary ``output_activation``.
     """
 
     def __init__(
@@ -115,7 +115,7 @@ class ConvolutionalAutoencoder(Autoencoder):
             # All intermediate layers should decode to
             # the same number of channels. The last decoder
             # should decode to whatever number of channels
-            # was specified, even if it's `None` (in which
+            # was specified, even if it's ``None`` (in which
             # case it will just be in_channels anyway)
             decode = in_channels if i else decode_channels
 

ml4gw/nn/resnet/resnet_1d.py

@@ -108,10 +108,10 @@ class BasicBlock(nn.Module):
 class Bottleneck(nn.Module):
     """
     Bottleneck blocks implement one extra convolution
-    compared to basic blocks. In this layers, the `planes`
-    parameter is generally meant to _downsize_ the number
+    compared to basic blocks. In this layers, the ``planes``
+    parameter is generally meant to **downsize** the number
     of feature maps first, which then get expanded out to
-    `planes * Bottleneck.expansion` feature maps at the
+    ``planes * Bottleneck.expansion`` feature maps at the
     output of the layer.
     """
 
@@ -192,9 +192,9 @@ class ResNet1D(nn.Module):
             A list representing the number of residual
             blocks to include in each "layer" of the
             network. Total layers (e.g. 50 in ResNet50)
-            is `2 + sum(layers) * factor`, where factor
-            is `2` for vanilla `ResNet` and `3` for
-            `BottleneckResNet`.
+            is ``2 + sum(layers) * factor``, where factor
+            is ``2`` for vanilla ``ResNet`` and ``3`` for
+            ``BottleneckResNet``.
         kernel_size:
             The size of the convolutional kernel to
             use in all residual layers. _NOT_ the size
@@ -211,19 +211,19 @@ class ResNet1D(nn.Module):
             connections between feature maps at subsequent
             layers rather than global. Generally won't
             need this to be >1, and wil raise an error if
-            >1 when using vanilla `ResNet`.
+            >1 when using vanilla ``ResNet``.
         width_per_group:
             Base width of each of the feature map groups,
             which is scaled up by the typical expansion
             factor at each layer of the network. Meaningless
-            for vanilla `ResNet`.
+            for vanilla ``ResNet``.
         stride_type:
             Whether to achieve downsampling on the time axis
             by strided or dilated convolutions for each layer.
-            If left as `None`, strided convolutions will be
-            used at each layer. Otherwise, `stride_type` should
-            be one element shorter than `layers` and indicate either
-            `stride` or `dilation` for each layer after the first.
+            If left as ``None``, strided convolutions will be
+            used at each layer. Otherwise, ``stride_type`` should
+            be one element shorter than ``layers`` and indicate either
+            ``stride`` or ``dilation`` for each layer after the first.
     """
 
     block = BasicBlock
@@ -316,7 +316,7 @@ class ResNet1D(nn.Module):
                 nn.init.kaiming_normal_(
                     m.weight, mode="fan_out", nonlinearity="relu"
                 )
-            elif isinstance(m, (nn.BatchNorm1d, nn.GroupNorm)):
+            elif isinstance(m, nn.BatchNorm1d):
                 nn.init.constant_(m.weight, 1)
                 nn.init.constant_(m.bias, 0)
 

ml4gw/nn/resnet/resnet_2d.py

@@ -105,10 +105,10 @@ class BasicBlock(nn.Module):
 class Bottleneck(nn.Module):
     """
     Bottleneck blocks implement one extra convolution
-    compared to basic blocks. In this layers, the `planes`
+    compared to basic blocks. In this layers, the ``planes``
     parameter is generally meant to _downsize_ the number
     of feature maps first, which then get expanded out to
-    `planes * Bottleneck.expansion` feature maps at the
+    ``planes * Bottleneck.expansion`` feature maps at the
     output of the layer.
     """
 
@@ -188,9 +188,9 @@ class ResNet2D(nn.Module):
             A list representing the number of residual
             blocks to include in each "layer" of the
             network. Total layers (e.g. 50 in ResNet50)
-            is `2 + sum(layers) * factor`, where factor
-            is `2` for vanilla `ResNet` and `3` for
-            `BottleneckResNet`.
+            is ``2 + sum(layers) * factor``, where factor
+            is ``2`` for vanilla ``ResNet`` and ``3`` for
+            ``BottleneckResNet``.
         kernel_size:
             The size of the convolutional kernel to
             use in all residual layers. _NOT_ the size
@@ -207,22 +207,22 @@ class ResNet2D(nn.Module):
             connections between feature maps at subsequent
             layers rather than global. Generally won't
             need this to be >1, and wil raise an error if
-            >1 when using vanilla `ResNet`.
+            >1 when using vanilla ``ResNet``.
         width_per_group:
             Base width of each of the feature map groups,
             which is scaled up by the typical expansion
             factor at each layer of the network. Meaningless
-            for vanilla `ResNet`.
+            for vanilla ``ResNet``.
         stride_type:
             Whether to achieve downsampling on the time axis
             by strided or dilated convolutions for each layer.
-            If left as `None`, strided convolutions will be
-            used at each layer. Otherwise, `stride_type` should
-            be one element shorter than `layers` and indicate either
-            `stride` or `dilation` for each layer after the first.
+            If left as ``None``, strided convolutions will be
+            used at each layer. Otherwise, ``stride_type`` should
+            be one element shorter than ``layers`` and indicate either
+            ``stride`` or ``dilation`` for each layer after the first.
         norm_groups:
             The number of groups to use in GroupNorm layers
-            throughout the model. If left as `-1`, the number
+            throughout the model. If left as ``-1``, the number
             of groups will be equal to the number of channels,
             making this equilavent to LayerNorm
     """

ml4gw/nn/streaming/online_average.py

@@ -11,7 +11,7 @@ class OnlineAverager(torch.nn.Module):
     """
     Module for performing stateful online averaging of
     batches of overlapping timeseries. At present, the
-    first `num_updates` predictions produced by this
+    first ``num_updates`` predictions produced by this
     model will underestimate the true average.
 
     Args:

ml4gw/nn/streaming/snapshotter.py

@@ -12,24 +12,24 @@ class Snapshotter(torch.nn.Module):
     Model for converting streaming state updates into
     a batch of overlapping snaphots of a multichannel
     timeseries. Can support multiple timeseries in a
-    single state update via the `channels_per_snapshot`
+    single state update via the ``channels_per_snapshot``
     kwarg.
 
     Specifically, maps tensors of shape
-    `(num_channels, batch_size * stride_size)` to a tensor
-    of shape `(batch_size, num_channels, snapshot_size)`.
-    If `channels_per_snapshot` is specified, it will return
-    `len(channels_per_snapshot)` tensors of this shape,
+    ``(num_channels, batch_size * stride_size)`` to a tensor
+    of shape ``(batch_size, num_channels, snapshot_size)``.
+    If ``channels_per_snapshot`` is specified, it will return
+    ``len(channels_per_snapshot)`` tensors of this shape,
     with the channel dimension replaced by the corresponding
-    value of `channels_per_snapshot`. The last tensor returned
+    value of ``channels_per_snapshot``. The last tensor returned
     at call time will be the current state that can be passed
-    to the next `forward` call.
+    to the next ``forward`` call.
 
     Args:
         num_channels:
             Number of channels in the timeseries. If
-            `channels_per_snapshot` is not `None`,
-            this should be equal to `sum(channels_per_snapshot)`.
+            ``channels_per_snapshot`` is not ``None``,
+            this should be equal to ``sum(channels_per_snapshot)``.
         snapshot_size:
             The size of the output snapshot windows in
             number of samples
@@ -39,17 +39,17 @@ class Snapshotter(torch.nn.Module):
         batch_size:
             The number of snapshots to produce at each
             update. The last dimension of the input
-            tensor should have size `batch_size * stride_size`.
+            tensor should have size ``batch_size * stride_size``.
         channels_per_snapshot:
             How to split up the channels in the timeseries
-            for different tensors. If left as `None`, all
+            for different tensors. If left as ``None``, all
             the channels will be returned in a single tensor.
             Otherwise, the channels will be split up into
-            `len(channels_per_snapshot)` tensors, with each
+            ``len(channels_per_snapshot)`` tensors, with each
             tensor's channel dimension being equal to the
-            corresponding value in `channels_per_snapshot`.
+            corresponding value in ``channels_per_snapshot``.
             Therefore, if specified, these values should
-            add up to `num_channels`.
+            add up to ``num_channels``.
     """
 
     def __init__(