Update test to reflect changes

NKI-AI · Apr 8, 2024 · 99d4b7a · 99d4b7a
1 parent 6062709
commit 99d4b7a
Showing 1 changed file with 24 additions and 2 deletions.
diff --git a/tests/tests_common/test_subsample.py b/tests/tests_common/test_subsample.py
@@ -22,8 +22,6 @@
         FastMRIMagicMaskFunc,
         Gaussian1DMaskFunc,
         Gaussian2DMaskFunc,
-        RadialMaskFunc,
-        SpiralMaskFunc,
         VariableDensityPoissonMaskFunc,
     ],
 )
@@ -44,6 +42,30 @@ def test_mask_reuse(mask_func, center_fracs, accelerations, batch_size, dim):
     assert torch.all(mask2 == mask3)
 
 
+@pytest.mark.parametrize(
+    "mask_func",
+    [
+        RadialMaskFunc,
+        SpiralMaskFunc,
+    ],
+)
+@pytest.mark.parametrize(
+    "accelerations, batch_size, dim",
+    [
+        ([4], 4, 320),
+        ([4, 8], 2, 368),
+    ],
+)
+def test_mask_reuse_circus(mask_func, accelerations, batch_size, dim):
+    mask_func = mask_func(accelerations=accelerations)
+    shape = (batch_size, dim, dim, 2)
+    mask1 = mask_func(shape, seed=123)
+    mask2 = mask_func(shape, seed=123)
+    mask3 = mask_func(shape, seed=123)
+    assert torch.all(mask1 == mask2)
+    assert torch.all(mask2 == mask3)
+
+
 @pytest.mark.parametrize(
     "mask_func",
     [