Add crop_pad_3d with tests (#679)

* add initial crop_pad_3d

* crop 3d test

* square crop 3d

* fixed crop3d test

* add padding

* pad 3d tests

* extra tests

* fix comment

* typo

Author: chris-langfield

src/aspire/utils/__init__.py

@@ -1,6 +1,7 @@
 from .coor_trans import (  # isort:skip
     common_line_from_rots,
     crop_pad_2d,
+    crop_pad_3d,
     get_aligned_rotations,
     get_rots_mse,
     grid_1d,

src/aspire/utils/coor_trans.py

@@ -333,3 +333,29 @@ def crop_pad_2d(im, size, fill_value=0):
     else:
         # target size is between mat_x and mat_y
         raise ValueError("Cannot crop and pad an image at the same time.")
+
+
+def crop_pad_3d(im, size, fill_value=0):
+    im_y, im_x, im_z = im.shape
+    # shift terms
+    start_x = math.floor(im_x / 2) - math.floor(size / 2)
+    start_y = math.floor(im_y / 2) - math.floor(size / 2)
+    start_z = math.floor(im_z / 2) - math.floor(size / 2)
+
+    # cropping
+    if size <= min(im_y, im_x, im_z):
+        return im[
+            start_y : start_y + size, start_x : start_x + size, start_z : start_z + size
+        ]
+    # padding
+    elif size >= max(im_y, im_x, im_z):
+        to_return = fill_value * np.ones((size, size, size), dtype=im.dtype)
+        to_return[
+            -start_y : im_y - start_y,
+            -start_x : im_x - start_x,
+            -start_z : im_z - start_z,
+        ] = im
+        return to_return
+    else:
+        # target size is between min and max of (im_y, im_x, im_z)
+        raise ValueError("Cannot crop and pad a volume at the same time.")

tests/test_coor_trans.py

@@ -6,6 +6,7 @@
 from aspire.utils import (
     Rotation,
     crop_pad_2d,
+    crop_pad_3d,
     get_aligned_rotations,
     grid_2d,
     grid_3d,
@@ -115,6 +116,36 @@ def testSquareCrop2D(self):
         test_a = np.diag(np.arange(8))
         self.assertTrue(np.array_equal(test_a, crop_pad_2d(a, 8)))
 
+    def testSquareCrop3D(self):
+        # even to even
+        a = np.zeros((8, 8, 8))
+        # pad it with the parts that will be cropped off from a 10x10x10
+        a = np.pad(a, ((1, 1), (1, 1), (1, 1)), "constant", constant_values=1)
+        # after cropping
+        test_a = np.zeros((8, 8, 8))
+        self.assertTrue(np.array_equal(crop_pad_3d(a, 8), test_a))
+
+        # even to odd
+        a = np.zeros((7, 7, 7))
+        # pad it with the parts that will be cropped off from a 10x10x10
+        a = np.pad(a, ((2, 1), (2, 1), (2, 1)), "constant", constant_values=1)
+        test_a = np.zeros((7, 7, 7))
+        self.assertTrue(np.array_equal(crop_pad_3d(a, 7), test_a))
+
+        # odd to odd
+        a = np.zeros((7, 7, 7))
+        # pad it with the parts that will be cropped off from a 9x9x9
+        a = np.pad(a, ((1, 1), (1, 1), (1, 1)), "constant", constant_values=1)
+        test_a = np.zeros((7, 7, 7))
+        self.assertTrue(np.array_equal(crop_pad_3d(a, 7), test_a))
+
+        # odd to even
+        a = np.zeros((8, 8, 8))
+        # pad it with the parts that will be cropped off from 11x11x11
+        a = np.pad(a, ((1, 2), (1, 2), (1, 2)), "constant", constant_values=1)
+        test_a = np.zeros((8, 8, 8))
+        self.assertTrue(np.array_equal(crop_pad_3d(a, 8), test_a))
+
     def testSquarePad2D(self):
         # Test even/odd cases based on the convention that the center of a sequence of length n
         # is (n+1)/2 if n is odd and n/2 + 1 if even.
@@ -151,6 +182,31 @@ def testSquarePad2D(self):
         test_a = np.diag([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
         self.assertTrue(np.array_equal(test_a, crop_pad_2d(a, 10)))
 
+    def testSquarePad3D(self):
+        # even to even
+        a = np.zeros((8, 8, 8))
+        # after padding to 10x10x10
+        test_a = np.pad(a, ((1, 1), (1, 1), (1, 1)), "constant", constant_values=1)
+        self.assertTrue(np.array_equal(crop_pad_3d(a, 10, fill_value=1), test_a))
+
+        # even to odd
+        a = np.zeros((8, 8, 8))
+        # after padding to 11x11x11
+        test_a = np.pad(a, ((1, 2), (1, 2), (1, 2)), "constant", constant_values=1)
+        self.assertTrue(np.array_equal(crop_pad_3d(a, 11, fill_value=1), test_a))
+
+        # odd to odd
+        a = np.zeros((7, 7, 7))
+        # after padding to 9x9x9
+        test_a = np.pad(a, ((1, 1), (1, 1), (1, 1)), "constant", constant_values=1)
+        self.assertTrue(np.array_equal(crop_pad_3d(a, 9, fill_value=1), test_a))
+
+        # odd to even
+        a = np.zeros((7, 7, 7))
+        # after padding to 10x10x10
+        test_a = np.pad(a, ((2, 1), (2, 1), (2, 1)), "constant", constant_values=1)
+        self.assertTrue(np.array_equal(crop_pad_3d(a, 10, fill_value=1), test_a))
+
     def testRectCrop2D(self):
         # Additional sanity checks for rectangular cropping case
 
@@ -240,10 +296,17 @@ def testRectPad2D(self):
     def testCropPad2DError(self):
         with self.assertRaises(ValueError) as e:
             _ = crop_pad_2d(np.zeros((6, 10)), 8)
-            self.assertTrue(
+            self.assertEqual(
                 "Cannot crop and pad an image at the same time.", str(e.exception)
             )
 
+    def testCropPad3DError(self):
+        with self.assertRaises(ValueError) as e:
+            _ = crop_pad_3d(np.zeros((6, 8, 10)), 8)
+            self.assertEqual(
+                "Cannot crop and pad a volume at the same time.", str(e.exception)
+            )
+
     def testCrop2DDtype(self):
         # crop_pad_2d must return an array of the same dtype it was given
         # in particular, because the method is used for Fourier downsampling
@@ -252,10 +315,23 @@ def testCrop2DDtype(self):
             crop_pad_2d(np.eye(10).astype("complex"), 5).dtype, np.dtype("complex128")
         )
 
+    def testCrop3DDtype(self):
+        self.assertEqual(
+            crop_pad_3d(np.ones((8, 8, 8)).astype("complex"), 5).dtype,
+            np.dtype("complex128"),
+        )
+
     def testCrop2DFillValue(self):
         # make sure the fill value is as expected
         # we are padding from an odd to an even dimension
         # so the padded column is added to the left
         a = np.ones((4, 3))
         b = crop_pad_2d(a, 4, fill_value=-1)
         self.assertTrue(np.array_equal(b[:, 0], np.array([-1, -1, -1, -1])))
+
+    def testCrop3DFillValue(self):
+        # make sure the fill value is expected. Since we are padding from odd to even
+        # the padded side is added to the 0-end of dimension 3
+        a = np.ones((4, 4, 3))
+        b = crop_pad_3d(a, 4, fill_value=-1)
+        self.assertTrue(np.array_equal(b[:, :, 0], -1 * np.ones((4, 4))))