ml4gw 0.7.0__py3-none-any.whl → 0.7.2__py3-none-any.whl

ml4gw/gw.py

@@ -16,8 +16,6 @@ import torch
 from jaxtyping import Float
 from torch import Tensor
 
-from ml4gw.utils.interferometer import InterferometerGeometry
-
 from .constants import C
 from .types import (
     BatchTensor,
@@ -28,6 +26,7 @@ from .types import (
     VectorGeometry,
     WaveformTensor,
 )
+from .utils.interferometer import InterferometerGeometry
 
 
 def outer(x: VectorGeometry, y: VectorGeometry) -> TensorGeometry:
@@ -285,6 +284,7 @@ def compute_ifo_snr(
     psd: PSDTensor,
     sample_rate: float,
     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
 
@@ -300,7 +300,8 @@ def compute_ifo_snr(
         {S_n^{(j)}(f)}df$$
 
     Where $f_{\text{min}}$ is a minimum frequency denoted
-    by `highpass`, `f_{\text{max}}` is the Nyquist frequency
+    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;
@@ -328,8 +329,15 @@ 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 `sample_rate / 2`
+            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
+            the Nyquist freqyency will contribute to the SNR calculation.
     Returns:
         Batch of SNRs computed for each interferometer
     """
@@ -346,7 +354,7 @@ def compute_ifo_snr(
     integrand = fft / (psd**0.5)
     integrand = integrand.type(torch.float32) ** 2
 
-    # mask out low frequency components if a critical
+    # mask out frequency components if a critical
     # frequency or frequency mask was provided
     if highpass is not None:
         if not isinstance(highpass, torch.Tensor):
@@ -354,12 +362,24 @@ def compute_ifo_snr(
             highpass = freqs >= highpass
         elif len(highpass) != integrand.shape[-1]:
             raise ValueError(
-                "Can't apply highpass filter mask with {} frequecy bins"
+                "Can't apply highpass filter mask with {} frequency bins"
                 "to signal fft with {} frequency bins".format(
                     len(highpass), integrand.shape[-1]
                 )
             )
         integrand *= highpass.to(integrand.device)
+    if lowpass is not None:
+        if not isinstance(lowpass, torch.Tensor):
+            freqs = torch.fft.rfftfreq(responses.shape[-1], 1 / sample_rate)
+            lowpass = freqs < lowpass
+        elif len(lowpass) != integrand.shape[-1]:
+            raise ValueError(
+                "Can't apply lowpass filter mask with {} frequency bins"
+                "to signal fft with {} frequency bins".format(
+                    len(lowpass), integrand.shape[-1]
+                )
+            )
+        integrand *= lowpass.to(integrand.device)
 
     # sum over the desired frequency range and multiply
     # by df to turn it into an integration (and get
@@ -386,6 +406,7 @@ def compute_network_snr(
     psd: PSDTensor,
     sample_rate: float,
     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
@@ -422,10 +443,17 @@ def compute_network_snr(
             If a float, it will be used to compute such a mask. If
             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
+            the Nyquist freqyency will contribute to the SNR calculation.
     Returns:
         Batch of SNRs for each waveform across the interferometer network
     """
-    snrs = compute_ifo_snr(responses, psd, sample_rate, highpass)
+    snrs = compute_ifo_snr(responses, psd, sample_rate, highpass, lowpass)
     snrs = snrs**2
     return snrs.sum(axis=-1) ** 0.5
 
@@ -436,6 +464,7 @@ def reweight_snrs(
     psd: PSDTensor,
     sample_rate: float,
     highpass: Union[float, Float[Tensor, " frequency"], None] = None,
+    lowpass: Union[float, Float[Tensor, " frequency"], None] = None,
 ) -> WaveformTensor:
     """Scale interferometer responses such that they have a desired SNR
 
@@ -466,10 +495,17 @@ def reweight_snrs(
             If a float, it will be used to compute such a mask. If
             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
+            the Nyquist freqyency will contribute to the SNR calculation.
     Returns:
         Rescaled interferometer responses
     """
 
-    snrs = compute_network_snr(responses, psd, sample_rate, highpass)
+    snrs = compute_network_snr(responses, psd, sample_rate, highpass, lowpass)
     weights = target_snrs / snrs
     return responses * weights[:, None, None]

ml4gw/spectral.py

@@ -382,7 +382,7 @@ def truncate_inverse_power_spectrum(
             as `None`, no lowpass filtering will be applied.
     Returns:
         The PSD with its time domain response truncated
-            to `fduration` and any highpassed frequencies
+            to `fduration` and any filtered frequencies
             tapered.
     """
 

ml4gw/transforms/snr_rescaler.py

@@ -14,10 +14,10 @@ class SnrRescaler(FittableSpectralTransform):
         sample_rate: float,
         waveform_duration: float,
         highpass: Optional[float] = None,
+        lowpass: Optional[float] = None,
         dtype: torch.dtype = torch.float32,
     ) -> None:
         super().__init__()
-        self.highpass = highpass
         self.sample_rate = sample_rate
         self.num_channels = num_channels
 
@@ -29,9 +29,18 @@ class SnrRescaler(FittableSpectralTransform):
 
         if highpass is not None:
             freqs = torch.fft.rfftfreq(waveform_size, 1 / sample_rate)
-            self.register_buffer("mask", freqs >= highpass, persistent=False)
+            self.register_buffer(
+                "highpass_mask", freqs >= highpass, persistent=False
+            )
+        else:
+            self.highpass_mask = None
+        if lowpass is not None:
+            freqs = torch.fft.rfftfreq(waveform_size, 1 / sample_rate)
+            self.register_buffer(
+                "lowpass_mask", freqs < lowpass, persistent=False
+            )
         else:
-            self.mask = None
+            self.lowpass_mask = None
 
     def fit(
         self,
@@ -63,7 +72,11 @@ class SnrRescaler(FittableSpectralTransform):
         target_snrs: Optional[BatchTensor] = None,
     ):
         snrs = compute_network_snr(
-            responses, self.background, self.sample_rate, self.mask
+            responses,
+            self.background,
+            self.sample_rate,
+            self.highpass_mask,
+            self.lowpass_mask,
         )
         if target_snrs is None:
             idx = torch.randperm(len(snrs))

ml4gw/waveforms/generator.py

@@ -224,7 +224,7 @@ class TimeDomainCBCWaveformGenerator(torch.nn.Module):
         k1s = torch.round(f_min / df)
 
         num_freqs = frequencies.size(0)
-        frequency_indices = torch.arange(num_freqs)
+        frequency_indices = torch.arange(num_freqs, device=device)
         taper_mask = frequency_indices <= k1s[:, None]
         taper_mask &= frequency_indices >= k0s[:, None]
 
@@ -253,7 +253,7 @@ class TimeDomainCBCWaveformGenerator(torch.nn.Module):
         # that will translate the coalescense time such that it is `right_pad`
         # seconds from the right edge of the window
         tshift = round(self.right_pad * self.sample_rate) / self.sample_rate
-        kvals = torch.arange(num_freqs)
+        kvals = torch.arange(num_freqs, device=device)
         phase_shift = torch.exp(1j * 2 * torch.pi * df * tshift * kvals)
 
         hc_spectrum *= phase_shift

{ml4gw-0.7.0.dist-info → ml4gw-0.7.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: ml4gw
-Version: 0.7.0
+Version: 0.7.2
 Summary: Tools for training torch models on gravitational wave data
 Author: Alec Gunny
 Author-email: alec.gunny@ligo.org

{ml4gw-0.7.0.dist-info → ml4gw-0.7.2.dist-info}/RECORD

@@ -6,7 +6,7 @@ ml4gw/dataloading/chunked_dataset.py,sha256=j96Rd67cRpsvotR_dzgfbrqxcoGDWnTV5cmf
 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=aUPSXgwyqJUBGGaKtUa-O3qkSbRYZwhhXIlkhvjgJgI,17684
+ml4gw/gw.py,sha256=0I9MhoHWksWG9a5EUI0GkHD1skuOXiaQgSgxNKYXCxE,19778
 ml4gw/nn/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ml4gw/nn/autoencoder/__init__.py,sha256=ZaT1XhJTHpMuPQqu5E__Jezeh9uwtjcXlT7IZ18byq4,161
 ml4gw/nn/autoencoder/base.py,sha256=eSWrDdpblI609oqa7RDSvZiY3YcV8WfhTioWKFn_7eE,3205
@@ -20,13 +20,13 @@ ml4gw/nn/resnet/resnet_2d.py,sha256=fVzYRuO0xR9yGjjQExv30mouokvupOAW-Kfdbs5WYDA,
 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=lLpnho02i-0zPSi96b0xOPEIgQMnBrmO8JiV1KvPGEw,19811
+ml4gw/spectral.py,sha256=rnxd1ObPjyQMAu3D83_sw2lEEHZF7f87YQBV_pxHLxM,19809
 ml4gw/transforms/__init__.py,sha256=OaTQJD4GFkDkcxt0DIwt2AzeEcv9t21ciKXxQnqDiuI,447
 ml4gw/transforms/iirfilter.py,sha256=RwgC3DWgYmBnHe7bYjvr9njM1WrRZ9EjBJsZNmtOY8s,3186
 ml4gw/transforms/pearson.py,sha256=CM9FTRxI4384-36FIaJFOcMZwsA7BkgberToJkMU1PA,3227
 ml4gw/transforms/qtransform.py,sha256=5S9y3PxkOmqMAarQmme0Tiy58vRvberpqhg6IeyDJLI,20675
 ml4gw/transforms/scaler.py,sha256=K5mp4w2zGZbpH1AcBUfpQS4n3aVSNzkaGWXedwk2LXs,2508
-ml4gw/transforms/snr_rescaler.py,sha256=XHKTeJXM3F_VOmjWOZetQuVZ6PMum8pEBPaOVbS16-w,2327
+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
@@ -48,8 +48,8 @@ ml4gw/waveforms/cbc/phenom_p.py,sha256=LBtGVUjBjROcYBPLldFnF6T1jZV6ZyuZEnkn9-oTK
 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=yCMsL55HRYHh6AVLWKA9Pmd6AoAyO7pe1VdrxzxxOzE,11975
-ml4gw-0.7.0.dist-info/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
-ml4gw-0.7.0.dist-info/METADATA,sha256=NsQW99R6S6MYDLoq_cFdgTK-lEZZBK8a-5S09RJIhZo,3904
-ml4gw-0.7.0.dist-info/WHEEL,sha256=IYZQI976HJqqOpQU6PHkJ8fb3tMNBFjg-Cn-pwAbaFM,88
-ml4gw-0.7.0.dist-info/RECORD,,
+ml4gw/waveforms/generator.py,sha256=i2lgaJzH5eA6gzc-bLQZYYEgEQ8OBLJgE9yNXU3FsKM,12005
+ml4gw-0.7.2.dist-info/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
+ml4gw-0.7.2.dist-info/METADATA,sha256=enY2IXiMhiIKn4RBDS6JLEdq1xT9g7IDjV6pvHxBpes,3904
+ml4gw-0.7.2.dist-info/WHEEL,sha256=IYZQI976HJqqOpQU6PHkJ8fb3tMNBFjg-Cn-pwAbaFM,88
+ml4gw-0.7.2.dist-info/RECORD,,