diff --git a/reconstructSPI/iterative_refinement/expectation_maximization.py b/reconstructSPI/iterative_refinement/expectation_maximization.py index 4571e60..ebaa780 100644 --- a/reconstructSPI/iterative_refinement/expectation_maximization.py +++ b/reconstructSPI/iterative_refinement/expectation_maximization.py @@ -141,6 +141,11 @@ 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, @@ -148,6 +153,11 @@ def iterative_refinement(self, wiener_small_number=0.01, count_norm_const=1): ) = 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] diff --git a/tests/test_expectation_maximization.py b/tests/test_expectation_maximization.py index 52af1fc..598179f 100644 --- a/tests/test_expectation_maximization.py +++ b/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)