Get tests running

ML4GW · Feb 27, 2025 · 446e708 · 446e708
1 parent 9f7c9e5
commit 446e708
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 6 deletions.
diff --git a/ml4gw/transforms/spectrogram.py b/ml4gw/transforms/spectrogram.py
@@ -104,7 +104,8 @@ def __init__(
         self.register_buffer("time_idxs", time_idxs)
 
     def _check_and_format_kwargs(self, kwargs: Dict[str, List]) -> List:
-        lengths = sorted((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(

diff --git a/ml4gw/transforms/spline_interpolation.py b/ml4gw/transforms/spline_interpolation.py
@@ -83,12 +83,14 @@ def __init__(
         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
         self.sy = sy
         self.register_buffer("x_in", x_in)
-        self.register_buffer("y_in", y_in or Tensor([1]))
+        self.register_buffer("y_in", y_in)
         self.register_buffer("x_out", x_out)
         self.register_buffer("y_out", y_out)
 

diff --git a/tests/dataloading/test_chunked_dataset.py b/tests/dataloading/test_chunked_dataset.py
@@ -66,6 +66,7 @@ def dataset(
         )
 
     def test_iter(self, dataset, coincident):
+        length = 0
         for x in dataset:
             assert x.shape == (8, 2, 192)
             if coincident:
@@ -76,5 +77,6 @@ def test_iter(self, dataset, coincident):
             diffs = torch.diff(x[:, 0], axis=-1)
             expected = torch.ones_like(diffs)
             torch.testing.assert_close(diffs, expected, rtol=0, atol=0)
+            length += 1
 
-        assert len(dataset) == 42
+        assert length == 42
diff --git a/tests/test_gw.py b/tests/test_gw.py
@@ -428,7 +428,9 @@ def test_reweight_snrs(_get_waveforms_from_lalsimulation):
     )
     # mutate data in the hp timeseries, and recompute snr using LAL
     hp.data.data = reweighted_response[..., 0, :].numpy().flatten()
+    ligo_snr = lalsimulation.MeasureSNR(hp, psd_1, 1, sample_rate / 2)
+    hp.data.data = reweighted_response[..., 1, :].numpy().flatten()
+    virgo_snr = lalsimulation.MeasureSNR(hp, psd_1, 1, sample_rate / 2)
+    network_snr = (ligo_snr**2 + virgo_snr**2) ** 0.5
 
-    assert lalsimulation.MeasureSNR(hp, psd_1, 1, 100) == pytest.approx(
-        target_network_snr.numpy()
-    )
+    assert network_snr == pytest.approx(target_network_snr.numpy(), rel=1e-1)