ml4gw 0.5.0__py3-none-any.whl → 0.6.0__py3-none-any.whl

ml4gw/spectral.py

@@ -12,14 +12,18 @@ https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.csd.html
 from typing import Optional, Union
 
 import torch
-from torchtyping import TensorType
+from jaxtyping import Float
+from torch import Tensor
 
-from ml4gw import types
+from ml4gw.types import (
+    FrequencySeries1to3d,
+    PSDTensor,
+    TimeSeries1to3d,
+    WaveformTensor,
+)
 
-time = None
 
-
-def median(x, axis):
+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
@@ -33,7 +37,7 @@ def median(x, axis):
 
 
 def _validate_shapes(
-    x: torch.Tensor, nperseg: int, y: Optional[torch.Tensor] = None
+    x: Tensor, nperseg: int, y: Optional[Tensor] = None
 ) -> None:
     if x.shape[-1] < nperseg:
         raise ValueError(
@@ -83,14 +87,14 @@ def _validate_shapes(
 
 
 def fast_spectral_density(
-    x: torch.Tensor,
+    x: TimeSeries1to3d,
     nperseg: int,
     nstride: int,
-    window: torch.Tensor,
-    scale: torch.Tensor,
+    window: Float[Tensor, " {nperseg//2+1}"],
+    scale: float,
     average: str = "median",
-    y: Optional[torch.Tensor] = None,
-) -> torch.Tensor:
+    y: Optional[TimeSeries1to3d] = None,
+) -> FrequencySeries1to3d:
     """
     Compute the power spectral density of a multichannel
     timeseries or a batch of multichannel timeseries, or
@@ -107,9 +111,9 @@ def fast_spectral_density(
             The timeseries tensor whose power spectral density
             to compute, or for cross spectral density the
             timeseries whose fft will be conjugated. Can have
-            shape either
-            `(batch_size, num_channels, length * sample_rate)`
-            or `(num_channels, length * sample_rate)`.
+            shape `(batch_size, num_channels, length * sample_rate)`,
+            `(num_channels, length * sample_rate)`, or
+            `(length * sample_rate)`.
         nperseg:
             Number of samples included in each FFT window
         nstride:
@@ -150,7 +154,7 @@ def fast_spectral_density(
             channel in `x` across _all_ of `x`'s batch elements.
     Returns:
         Tensor of power spectral densities of `x` or its cross spectral
-            density with the timeseries in `y`.
+        density with the timeseries in `y`.
     """
 
     _validate_shapes(x, nperseg, y)
@@ -240,17 +244,16 @@ def fast_spectral_density(
 
 
 def spectral_density(
-    x: torch.Tensor,
+    x: TimeSeries1to3d,
     nperseg: int,
     nstride: int,
-    window: torch.Tensor,
-    scale: torch.Tensor,
+    window: Float[Tensor, " {nperseg//2+1}"],
+    scale: float,
     average: str = "median",
-) -> torch.Tensor:
+) -> FrequencySeries1to3d:
     """
     Compute the power spectral density of a multichannel
-    timeseries or a batch of multichannel timeseries, or
-    the cross power spectral density of two such timeseries.
+    timeseries or a batch of multichannel timeseries.
     This implementation is exact for all frequency bins, but
     slower than the fast implementation.
 
@@ -259,9 +262,9 @@ def spectral_density(
             The timeseries tensor whose power spectral density
             to compute, or for cross spectral density the
             timeseries whose fft will be conjugated. Can have
-            shape either
-            `(batch_size, num_channels, length * sample_rate)`
-            or `(num_channels, length * sample_rate)`.
+            shape `(batch_size, num_channels, length * sample_rate)`,
+            `(num_channels, length * sample_rate)`, or
+            `(length * sample_rate)`.
         nperseg:
             Number of samples included in each FFT window
         nstride:
@@ -336,11 +339,11 @@ def spectral_density(
 
 
 def truncate_inverse_power_spectrum(
-    psd: types.PSDTensor,
-    fduration: Union[TensorType["time"], float],
+    psd: PSDTensor,
+    fduration: Union[Float[Tensor, " time"], float],
     sample_rate: float,
     highpass: Optional[float] = None,
-) -> types.PSDTensor:
+) -> PSDTensor:
     """
     Truncate the length of the time domain response
     of a whitening filter built using the specified
@@ -399,7 +402,7 @@ def truncate_inverse_power_spectrum(
     q = torch.fft.irfft(inv_asd, n=N, norm="forward", dim=-1)
 
     # taper the edges of the TD filter
-    if isinstance(fduration, torch.Tensor):
+    if isinstance(fduration, Tensor):
         pad = fduration.size(-1) // 2
         window = fduration
     else:
@@ -422,8 +425,8 @@ def truncate_inverse_power_spectrum(
 
 
 def normalize_by_psd(
-    X: types.WaveformTensor,
-    psd: types.PSDTensor,
+    X: WaveformTensor,
+    psd: PSDTensor,
     sample_rate: float,
     pad: int,
 ):
@@ -438,7 +441,7 @@ def normalize_by_psd(
 
     # convert back to the time domain and normalize
     # TODO: what's this normalization factor?
-    X = torch.fft.irfft(X_tilde, norm="forward", dim=-1)
+    X = torch.fft.irfft(X_tilde, n=X.shape[-1], norm="forward", dim=-1)
     X = X.float() / sample_rate**0.5
 
     # slice off corrupted data at edges of kernel
@@ -447,12 +450,12 @@ def normalize_by_psd(
 
 
 def whiten(
-    X: types.WaveformTensor,
-    psd: types.PSDTensor,
-    fduration: Union[TensorType["time"], float],
+    X: WaveformTensor,
+    psd: PSDTensor,
+    fduration: Union[Float[Tensor, " time"], float],
     sample_rate: float,
     highpass: Optional[float] = None,
-) -> types.WaveformTensor:
+) -> WaveformTensor:
     """
     Whiten a batch of timeseries using the specified
     background one-sided power spectral densities (PSDs),
@@ -460,7 +463,8 @@ def whiten(
     `fduration` and possibly to highpass filter.
 
     Args:
-        X: batch of multichannel timeseries to whiten
+        X:
+            batch of multichannel timeseries to whiten
         psd:
             PSDs use to whiten the data. The frequency
             response of the whitening filter will be roughly
@@ -496,7 +500,7 @@ def whiten(
 
     # figure out how much data we'll need to slice
     # off after whitening
-    if isinstance(fduration, torch.Tensor):
+    if isinstance(fduration, Tensor):
         pad = fduration.size(-1) // 2
     else:
         pad = int(fduration * sample_rate / 2)

ml4gw/transforms/__init__.py

@@ -4,5 +4,6 @@ from .scaler import ChannelWiseScaler
 from .snr_rescaler import SnrRescaler
 from .spectral import SpectralDensity
 from .spectrogram import MultiResolutionSpectrogram
+from .spline_interpolation import SplineInterpolate
 from .waveforms import WaveformProjector, WaveformSampler
 from .whitening import FixedWhiten, Whiten

ml4gw/transforms/pearson.py

@@ -1,5 +1,8 @@
 import torch
+from jaxtyping import Float
+from torch import Tensor
 
+from ml4gw.types import TimeSeries1to3d
 from ml4gw.utils.slicing import unfold_windows
 
 
@@ -40,7 +43,7 @@ class ShiftedPearsonCorrelation(torch.nn.Module):
         super().__init__()
         self.max_shift = max_shift
 
-    def _shape_checks(self, x: torch.Tensor, y: torch.Tensor):
+    def _shape_checks(self, x: TimeSeries1to3d, y: TimeSeries1to3d):
         if x.ndim > 3:
             raise ValueError(
                 "Tensor x can only have up to 3 dimensions "
@@ -61,8 +64,9 @@ class ShiftedPearsonCorrelation(torch.nn.Module):
                     )
                 )
 
-    # TODO: torchtyping annotate
-    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+    def forward(
+        self, x: TimeSeries1to3d, y: TimeSeries1to3d
+    ) -> Float[Tensor, "windows ..."]:
         self._shape_checks(x, y)
         dim = x.size(-1)
 

ml4gw/transforms/qtransform.py

@@ -1,8 +1,14 @@
 import math
-from typing import List, Optional, Tuple
+import warnings
+from typing import List, Tuple
 
 import torch
 import torch.nn.functional as F
+from jaxtyping import Float, Int
+from torch import Tensor
+
+from ml4gw.transforms.spline_interpolation import SplineInterpolate
+from ml4gw.types import FrequencySeries1to3d, TimeSeries1to3d, TimeSeries3d
 
 """
 All based on https://github.com/gwpy/gwpy/blob/v3.0.8/gwpy/signal/qtransform.py
@@ -34,7 +40,6 @@ class QTile(torch.nn.Module):
         mismatch:
             The maximum fractional mismatch between neighboring tiles
 
-
     """
 
     def __init__(
@@ -44,7 +49,7 @@ class QTile(torch.nn.Module):
         duration: float,
         sample_rate: float,
         mismatch: float,
-    ):
+    ) -> None:
         super().__init__()
         self.mismatch = mismatch
         self.q = q
@@ -63,18 +68,18 @@ class QTile(torch.nn.Module):
         self.register_buffer("indices", self.get_data_indices())
         self.register_buffer("window", self.get_window())
 
-    def ntiles(self):
+    def ntiles(self) -> int:
         """
         Number of tiles in this frequency row
         """
         tcum_mismatch = self.duration * 2 * torch.pi * self.frequency / self.q
         return int(2 ** torch.ceil(torch.log2(tcum_mismatch / self.deltam)))
 
-    def _get_indices(self):
+    def _get_indices(self) -> Int[Tensor, " windowsize"]:
         half = int((self.windowsize - 1) / 2)
         return torch.arange(-half, half + 1)
 
-    def get_window(self):
+    def get_window(self) -> Float[Tensor, " windowsize"]:
         """
         Generate the bi-square window for this row
         """
@@ -87,7 +92,7 @@ class QTile(torch.nn.Module):
         )
         return torch.Tensor((1 - xfrequencies**2) ** 2 * norm)
 
-    def get_data_indices(self):
+    def get_data_indices(self) -> Int[Tensor, " windowsize"]:
         """
         Get the index array of relevant frequencies for this row
         """
@@ -95,7 +100,11 @@ class QTile(torch.nn.Module):
             self._get_indices() + 1 + self.frequency * self.duration,
         ).type(torch.long)
 
-    def forward(self, fseries: torch.Tensor, norm: str = "median"):
+    def forward(
+        self,
+        fseries: FrequencySeries1to3d,
+        norm: str = "median",
+    ) -> TimeSeries1to3d:
         """
         Compute the transform for this row
 
@@ -138,7 +147,7 @@ class QTile(torch.nn.Module):
                 energy /= means
             else:
                 raise ValueError("Invalid normalisation %r" % norm)
-            return energy.type(torch.float32)
+            energy = energy.type(torch.float32)
         return energy
 
 
@@ -166,6 +175,19 @@ class SingleQTransform(torch.nn.Module):
             be chosen based on q, sample_rate, and duration
         mismatch:
             The maximum fractional mismatch between neighboring tiles
+        interpolation_method:
+            The method by which to interpolate each `QTile` to the specified
+            number of time and frequency bins. The acceptable values are
+            "bilinear", "bicubic", and "spline". The "bilinear" and "bicubic"
+            options will use PyTorch's built-in interpolation modes, while
+            "spline" will use the custom Torch-based implementation in
+            `ml4gw`, as PyTorch does not have spline-based intertpolation.
+            The "spline" mode is most similar to the results of GWpy's
+            Q-transform, which uses `scipy` to do spline interpolation.
+            However, it is also the slowest and most memory intensive due to
+            the matrix equation solving steps. Therefore, the default method
+            is "bicubic" as it produces the most similar results while
+            optimizing for computing performance.
     """
 
     def __init__(
@@ -176,7 +198,8 @@ class SingleQTransform(torch.nn.Module):
         q: float = 12,
         frange: List[float] = [0, torch.inf],
         mismatch: float = 0.2,
-    ):
+        interpolation_method: str = "bicubic",
+    ) -> None:
         super().__init__()
         self.q = q
         self.spectrogram_shape = spectrogram_shape
@@ -184,21 +207,88 @@ class SingleQTransform(torch.nn.Module):
         self.duration = duration
         self.mismatch = mismatch
 
+        # If q is too large, the minimum of the frange computed
+        # below will be larger than the maximum
+        max_q = torch.pi * duration * sample_rate / 50 - 11 ** (0.5)
+        if q >= max_q:
+            raise ValueError(
+                "The given q value is too large for the given duration and "
+                f"sample rate. The maximum allowable value is {max_q}"
+            )
+
+        if interpolation_method not in ["bilinear", "bicubic", "spline"]:
+            raise ValueError(
+                "Interpolation method must be either 'bilinear', 'bicubic', "
+                f"or 'spline'; got {interpolation_method}"
+            )
+        self.interpolation_method = interpolation_method
+
         qprime = self.q / 11 ** (1 / 2.0)
         if self.frange[0] <= 0:  # set non-zero lower frequency
             self.frange[0] = 50 * self.q / (2 * torch.pi * duration)
         if math.isinf(self.frange[1]):  # set non-infinite upper frequency
             self.frange[1] = sample_rate / 2 / (1 + 1 / qprime)
+
         self.freqs = self.get_freqs()
         self.qtile_transforms = torch.nn.ModuleList(
             [
-                QTile(self.q, freq, self.duration, sample_rate, self.mismatch)
+                QTile(
+                    q=self.q,
+                    frequency=freq,
+                    duration=self.duration,
+                    sample_rate=sample_rate,
+                    mismatch=self.mismatch,
+                )
                 for freq in self.freqs
             ]
         )
         self.qtiles = None
 
-    def get_freqs(self):
+        if self.interpolation_method == "spline":
+            self._set_up_spline_interp()
+
+    def _set_up_spline_interp(self):
+        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)))
+        idx = torch.arange(len(ntiles))
+        self.stack_idx = [idx[Tensor(ntiles) == n] for n in unique_ntiles]
+
+        t_out = torch.arange(
+            0, self.duration, self.duration / self.spectrogram_shape[1]
+        )
+        self.qtile_interpolators = torch.nn.ModuleList(
+            [
+                SplineInterpolate(
+                    kx=3,
+                    x_in=torch.arange(0, self.duration, self.duration / tiles),
+                    y_in=torch.arange(len(idx)),
+                    x_out=t_out,
+                    y_out=torch.arange(len(idx)),
+                )
+                for tiles, idx in zip(unique_ntiles, self.stack_idx)
+            ]
+        )
+
+        t_in = t_out
+        f_in = self.freqs
+        f_out = torch.logspace(
+            math.log10(self.frange[0]),
+            math.log10(self.frange[-1]),
+            self.spectrogram_shape[0],
+        )
+
+        self.interpolator = SplineInterpolate(
+            kx=3,
+            ky=3,
+            x_in=t_in,
+            y_in=f_in,
+            x_out=t_out,
+            y_out=f_out,
+        )
+
+    def get_freqs(self) -> Float[Tensor, " nfreq"]:
         """
         Calculate the frequencies that will be used in this transform.
         For each frequency, a `QTile` is created.
@@ -214,7 +304,8 @@ class SingleQTransform(torch.nn.Module):
 
         freq_base = math.exp(2 / ((2 + self.q**2) ** (1 / 2.0)) * fstep)
         freqs = torch.Tensor([freq_base ** (i + 0.5) for i in range(nfreq)])
-        freqs = (minf * freqs // fstepmin) * fstepmin
+        # Cast freqs to float64 to avoid off-by-ones from rounding
+        freqs = (minf * freqs.double() // fstepmin) * fstepmin
         return torch.unique(freqs)
 
     def get_max_energy(
@@ -262,7 +353,11 @@ class SingleQTransform(torch.nn.Module):
         if dimension == "batch":
             return torch.max(max_across_ft, dim=-1).values
 
-    def compute_qtiles(self, X: torch.Tensor, norm: str = "median"):
+    def compute_qtiles(
+        self,
+        X: TimeSeries1to3d,
+        norm: str = "median",
+    ) -> None:
         """
         Take the FFT of the input timeseries and calculate the transform
         for each `QTile`
@@ -272,36 +367,47 @@ class SingleQTransform(torch.nn.Module):
         X[..., 1:] *= 2
         self.qtiles = [qtile(X, norm) for qtile in self.qtile_transforms]
 
-    def interpolate(self, num_f_bins: int, num_t_bins: int):
-        """
-        Interpolate each `QTile` to the specified number of time and
-        frequency bins. Note that PyTorch does not have the same
-        interpolation methods that GWpy uses, and so the interpolated
-        spectrograms will be different even though the uninterpolated
-        values match. The `bicubic` interpolation method is used as
-        it seems to match GWpy most closely.
-        """
+    def interpolate(self) -> TimeSeries3d:
         if self.qtiles is None:
             raise RuntimeError(
                 "Q-tiles must first be computed with .compute_qtiles()"
             )
+        if self.interpolation_method == "spline":
+            qtiles = [
+                torch.stack([self.qtiles[i] for i in idx], dim=-2)
+                for idx in self.stack_idx
+            ]
+            time_interped = torch.cat(
+                [
+                    interpolator(qtile)
+                    for qtile, interpolator in zip(
+                        qtiles, self.qtile_interpolators
+                    )
+                ],
+                dim=-2,
+            )
+            return self.interpolator(time_interped)
+        num_f_bins, num_t_bins = self.spectrogram_shape
         resampled = [
             F.interpolate(
-                qtile[None], (qtile.shape[-2], num_t_bins), mode="bicubic"
+                qtile[None],
+                (qtile.shape[-2], num_t_bins),
+                mode=self.interpolation_method,
             )
             for qtile in self.qtiles
         ]
         resampled = torch.stack(resampled, dim=-2)
         resampled = F.interpolate(
-            resampled[0], (num_f_bins, num_t_bins), mode="bicubic"
+            resampled[0],
+            (num_f_bins, num_t_bins),
+            mode=self.interpolation_method,
         )
         return torch.squeeze(resampled)
 
     def forward(
         self,
-        X: torch.Tensor,
+        X: TimeSeries1to3d,
         norm: str = "median",
-        spectrogram_shape: Optional[Tuple[int, int]] = None,
     ):
         """
         Compute the Q-tiles and interpolate
@@ -315,24 +421,15 @@ class SingleQTransform(torch.nn.Module):
                 three-dimensional, axes will be added during Q-tile
                 computation.
             norm:
-                The method of interpolation used by each QTile
-            spectrogram_shape:
-                The shape of the interpolated spectrogram, specified as
-                `(num_f_bins, num_t_bins)`. Because the
-                frequency spacing of the Q-tiles is in log-space, the frequency
-                interpolation is log-spaced as well. If not given, the shape
-                used to initialize the transform will be used.
+                The method of normalization used by each QTile
 
         Returns:
             The interpolated Q-transform for the batch of data. Output will
             have one more dimension than the input
         """
 
-        if spectrogram_shape is None:
-            spectrogram_shape = self.spectrogram_shape
-        num_f_bins, num_t_bins = spectrogram_shape
         self.compute_qtiles(X, norm)
-        return self.interpolate(num_f_bins, num_t_bins)
+        return self.interpolate()
 
 
 class QScan(torch.nn.Module):
@@ -370,14 +467,22 @@ class QScan(torch.nn.Module):
         spectrogram_shape: Tuple[int, int],
         qrange: List[float] = [4, 64],
         frange: List[float] = [0, torch.inf],
+        interpolation_method="bicubic",
         mismatch: float = 0.2,
-    ):
+    ) -> None:
         super().__init__()
         self.qrange = qrange
         self.mismatch = mismatch
-        self.qs = self.get_qs()
         self.frange = frange
         self.spectrogram_shape = spectrogram_shape
+        max_q = torch.pi * duration * sample_rate / 50 - 11 ** (0.5)
+        self.qs = self.get_qs()
+        if self.qs[-1] >= max_q:
+            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."
+            )
 
         # Deliberately doing something different from GWpy here.
         # Their final frange is the intersection of the frange
@@ -391,13 +496,15 @@ class QScan(torch.nn.Module):
                     spectrogram_shape=spectrogram_shape,
                     q=q,
                     frange=self.frange.copy(),
+                    interpolation_method=interpolation_method,
                     mismatch=self.mismatch,
                 )
                 for q in self.qs
+                if q < max_q
             ]
         )
 
-    def get_qs(self):
+    def get_qs(self) -> List[float]:
         """
         Determine the values of Q to try for the set of Q-transforms
         """
@@ -409,14 +516,14 @@ class QScan(torch.nn.Module):
             self.qrange[0] * math.exp(2 ** (1 / 2.0) * dq * (i + 0.5))
             for i in range(nplanes)
         ]
+
         return qs
 
     def forward(
         self,
-        X: torch.Tensor,
+        X: TimeSeries1to3d,
         fsearch_range: List[float] = None,
         norm: str = "median",
-        spectrogram_shape: Optional[Tuple[int, int]] = None,
     ):
         """
         Compute the set of QTiles for each Q transform and determine which
@@ -436,12 +543,6 @@ class QScan(torch.nn.Module):
                 for the maximum energy
             norm:
                 The method of interpolation used by each QTile
-            spectrogram_shape:
-                The shape of the interpolated spectrogram, specified as
-                `(num_f_bins, num_t_bins)`. Because the
-                frequency spacing of the Q-tiles is in log-space, the frequency
-                interpolation is log-spaced as well. If not given, the shape
-                used to initialize the transform will be used.
 
         Returns:
             An interpolated Q-transform for the batch of data. Output will
@@ -457,7 +558,4 @@ class QScan(torch.nn.Module):
                 ]
             )
         )
-        if spectrogram_shape is None:
-            spectrogram_shape = self.spectrogram_shape
-        num_f_bins, num_t_bins = spectrogram_shape
-        return self.q_transforms[idx].interpolate(num_f_bins, num_t_bins)
+        return self.q_transforms[idx].interpolate()

ml4gw/transforms/scaler.py

@@ -1,6 +1,8 @@
 from typing import Optional
 
 import torch
+from jaxtyping import Float
+from torch import Tensor
 
 from ml4gw.transforms.transform import FittableTransform
 
@@ -34,7 +36,9 @@ class ChannelWiseScaler(FittableTransform):
         self.register_buffer("mean", mean)
         self.register_buffer("std", std)
 
-    def fit(self, X: torch.Tensor) -> None:
+    def fit(
+        self, X: Float[Tensor, "... time"], std_reg: Optional[float] = 0.0
+    ) -> None:
         """Fit the scaling parameters to a timeseries
 
         Computes the channel-wise mean and standard deviation
@@ -57,10 +61,12 @@ class ChannelWiseScaler(FittableTransform):
                 "Can't fit channel wise mean and standard deviation "
                 "from tensor of shape {}".format(X.shape)
             )
-
+        std += std_reg * torch.ones_like(std)
         super().build(mean=mean, std=std)
 
-    def forward(self, X: torch.Tensor, reverse: bool = False) -> torch.Tensor:
+    def forward(
+        self, X: Float[Tensor, "... time"], reverse: bool = False
+    ) -> Float[Tensor, "... time"]:
         if not reverse:
             return (X - self.mean) / self.std
         else:

ml4gw/transforms/snr_rescaler.py

@@ -2,8 +2,9 @@ from typing import Optional
 
 import torch
 
-from ml4gw import gw
+from ml4gw.gw import compute_network_snr
 from ml4gw.transforms.transform import FittableSpectralTransform
+from ml4gw.types import BatchTensor, TimeSeries2d, WaveformTensor
 
 
 class SnrRescaler(FittableSpectralTransform):
@@ -34,7 +35,7 @@ class SnrRescaler(FittableSpectralTransform):
 
     def fit(
         self,
-        *background: torch.Tensor,
+        *background: TimeSeries2d,
         fftlength: Optional[float] = None,
         overlap: Optional[float] = None,
     ):
@@ -58,10 +59,10 @@ class SnrRescaler(FittableSpectralTransform):
 
     def forward(
         self,
-        responses: gw.WaveformTensor,
-        target_snrs: Optional[gw.ScalarTensor] = None,
+        responses: WaveformTensor,
+        target_snrs: Optional[BatchTensor] = None,
     ):
-        snrs = gw.compute_network_snr(
+        snrs = compute_network_snr(
             responses, self.background, self.sample_rate, self.mask
         )
         if target_snrs is None: