Merge pull request #38 from compSPI/mask_tests

test_binary_mask_3d - circle inscribed in square

reconstructSPI/iterative_refinement/expectation_maximization.py

@@ -141,13 +141,23 @@ def iterative_refinement(self, wiener_small_number=0.01, count_norm_const=1):
                 ) = IterativeRefinement.compute_bayesian_weights(
                     particles_f_1[particle_idx], slices_conv_ctfs_1, sigma
                 )
+                print(
+                    "log z_norm_const_1={}, em_loss_1={}".format(
+                        z_norm_const_1, em_loss_1
+                    )
+                )
                 (
                     bayes_factors_2,
                     z_norm_const_2,
                     em_loss_2,
                 ) = IterativeRefinement.compute_bayesian_weights(
                     particles_f_2[particle_idx], slices_conv_ctfs_2, sigma
                 )
+                print(
+                    "log z_norm_const_2={}, em_loss_2={}".format(
+                        z_norm_const_2, em_loss_2
+                    )
+                )
 
                 for one_slice_idx in range(bayes_factors_1.shape[0]):
                     xyz = xyz_rotated[one_slice_idx]

tests/test_expectation_maximization.py

@@ -198,6 +198,70 @@ def test_compute_fsc(test_ir, n_pix):
     assert fsc_1.shape == (n_pix // 2,)
 
 
+def test_binary_mask_3d(test_ir):
+    """Test binary_mask_3d.
+
+    Tests the limit of infinite n_pix. Use high n_pix so good approx.
+    1. Sums shell through an axis, then converts to circle,
+    then checks if circle/square ratio agrees with largest
+    circle inscribed in square. Should be pi/4.
+
+    2. Make shells at sizes r and r/2 and check ratios of perimeter
+    of circle (mid slice) and surface area of sphere.
+
+    3. Make filled sphere of sizes r and r/2 and check ratio of volume.
+    """
+    n_pix = 512
+
+    center = (n_pix // 2, n_pix // 2, n_pix // 2)
+    radius = n_pix // 2
+    shape = (n_pix, n_pix, n_pix)
+    for fill in [True, False]:
+        mask = test_ir.binary_mask_3d(
+            center, radius, shape, fill=fill, shell_thickness=1
+        )
+
+        for axis in [0, 1, 2]:
+            circle = mask.sum(axis=axis) > 0
+            circle_to_square_ratio = circle.mean()
+            assert np.isclose(circle_to_square_ratio, np.pi / 4, atol=1e-3)
+
+    mask = test_ir.binary_mask_3d(center, radius, shape, fill=True, shell_thickness=1)
+    circle = mask[n_pix // 2]
+    circle_to_square_ratio = circle.mean()
+    assert np.isclose(circle_to_square_ratio, np.pi / 4, atol=1e-3)
+
+    r_half = radius / 2
+    for shell_thickness in [1, 2]:
+        mask_r = test_ir.binary_mask_3d(
+            center, radius, shape, fill=False, shell_thickness=1
+        )
+        mask_r_half = test_ir.binary_mask_3d(
+            center, r_half, shape, fill=False, shell_thickness=1
+        )
+        perimeter_ratio = mask_r[n_pix // 2].sum() / mask_r_half[n_pix // 2].sum()
+        assert np.isclose(2, perimeter_ratio, atol=0.1)
+        if shell_thickness == 1:
+            assert np.isclose(
+                mask_r[n_pix // 2].sum() / (2 * np.pi * radius), 1, atol=0.1
+            )
+            assert np.isclose(
+                mask_r_half[n_pix // 2].sum() / (2 * np.pi * r_half), 1, atol=0.1
+            )
+
+        surface_area_ratio = mask_r.sum() / mask_r_half.sum()
+        surface_area_ratio_analytic = (radius / r_half) ** 2
+        assert np.isclose(surface_area_ratio, surface_area_ratio_analytic, atol=0.1)
+
+    mask_r = test_ir.binary_mask_3d(center, radius, shape, fill=True, shell_thickness=1)
+    mask_r_half = test_ir.binary_mask_3d(
+        center, r_half, shape, fill=True, shell_thickness=1
+    )
+    volume_ratio = mask_r.sum() / mask_r_half.sum()
+    volume_ratio_analytic = (radius / r_half) ** 3
+    assert np.isclose(volume_ratio, volume_ratio_analytic, atol=0.005)
+
+
 def test_expand_1d_to_3d(test_ir, n_pix):
     """Test expansion of 1D array into spherical shell."""
     arr_1d = np.ones(n_pix // 2)