initial MicrographSource/Image extension to rect shapes

Author: garrettwrong

src/aspire/image/image.py

@@ -202,20 +202,19 @@ def __init__(self, data, pixel_size=None, dtype=None):
         else:
             self.dtype = np.dtype(dtype)
 
-        if not data.shape[-1] == data.shape[-2]:
-            raise ValueError("Only square ndarrays are supported.")
-
         self._data = data.astype(self.dtype, copy=False)
         self.ndim = self._data.ndim
         self.shape = self._data.shape
         self.stack_ndim = self._data.ndim - 2
         self.stack_shape = self._data.shape[:-2]
         self.n_images = np.prod(self.stack_shape)
-        self.resolution = self._data.shape[-1]
         self.pixel_size = None
         if pixel_size is not None:
             self.pixel_size = float(pixel_size)
 
+        self._is_square = data.shape[-1] == data.shape[-2]
+        self.resolution = self._data.shape[-1]  # XXXXX
+
         # Numpy interop
         # https://numpy.org/devdocs/user/basics.interoperability.html#the-array-interface-protocol
         self.__array_interface__ = self._data.__array_interface__
@@ -233,6 +232,12 @@ def project(self, angles):
         :return: Radon transform of the Image Stack.
         :rtype: Ndarray (stack size, number of angles, image resolution)
         """
+
+        if not self._is_square:
+            raise NotImplementedError(
+                "`Image.project` is not currently implemented for non-square images."
+            )
+
         # number of points to sample on radial line in polar grid
         n_points = self.resolution
         original_stack = self.stack_shape
@@ -308,19 +313,19 @@ def stack_reshape(self, *args):
         )
 
     def __add__(self, other):
-        if isinstance(other, Image):
+        if isinstance(other, self.__class__):
             other = other._data
 
         return self.__class__(self._data + other, pixel_size=self.pixel_size)
 
     def __sub__(self, other):
-        if isinstance(other, Image):
+        if isinstance(other, self.__class__):
             other = other._data
 
         return self.__class__(self._data - other, pixel_size=self.pixel_size)
 
     def __mul__(self, other):
-        if isinstance(other, Image):
+        if isinstance(other, self.__class__):
             other = other._data
 
         return self.__class__(self._data * other, pixel_size=self.pixel_size)
@@ -384,7 +389,7 @@ def __repr__(self):
             px_msg = f" with pixel_size={self.pixel_size} angstroms."
 
         msg = f"{self.n_images} {self.dtype} images arranged as a {self.stack_shape} stack"
-        msg += f" each of size {self.resolution}x{self.resolution}{px_msg}"
+        msg += f" each of size {self.shape[-2:]}{px_msg}"
         return msg
 
     def asnumpy(self):
@@ -441,6 +446,12 @@ def legacy_whiten(self, psd, delta):
             and which to set to zero. By default all `sqrt(psd)` values
             less than `delta` are zeroed out in the whitening filter.
         """
+
+        if not self._is_square:
+            raise NotImplementedError(
+                "`Image.legacy_whiten` is not currently implemented for non-square images."
+            )
+
         n = self.n_images
         L = self.resolution
         L_half = L // 2
@@ -577,6 +588,11 @@ def filter(self, filter):
         :param filter: An object of type `Filter`.
         :return: A new filtered `Image` object.
         """
+        if not self._is_square:
+            raise NotImplementedError(
+                "`Image.filter` is not currently implemented for non-square images."
+            )
+
         original_stack_shape = self.stack_shape
 
         im = self.stack_reshape(-1)
@@ -667,8 +683,8 @@ def _load_raw(filepath, dtype=None):
 
         return im, pixel_size
 
-    @staticmethod
-    def load(filepath, dtype=None):
+    @classmethod
+    def load(cls, filepath, dtype=None):
         """
         Load raw data from supported files.
 
@@ -683,7 +699,7 @@ def load(filepath, dtype=None):
         im, pixel_size = Image._load_raw(filepath, dtype=dtype)
 
         # Return as Image instance
-        return Image(im, pixel_size=pixel_size)
+        return cls(im, pixel_size=pixel_size)
 
     def _im_translate(self, shifts):
         """
@@ -699,6 +715,10 @@ def _im_translate(self, shifts):
             Alternatively, it can be a row vector of length 2, in which case the same shifts is applied to each image.
         :return: The images translated by the shifts, with periodic boundaries.
         """
+        if not self._is_square:
+            raise NotImplementedError(
+                "`Image._im_translate` is not currently implemented for non-square images."
+            )
 
         if shifts.ndim == 1:
             shifts = shifts[np.newaxis, :]
@@ -769,6 +789,10 @@ def backproject(self, rot_matrices, symmetry_group=None, zero_nyquist=True):
 
         :return: Volume instance corresonding to the backprojected images.
         """
+        if not self._is_square:
+            raise NotImplementedError(
+                "`Image.legacy_whiten` is not currently implemented for non-square images."
+            )
 
         if self.stack_ndim > 1:
             raise NotImplementedError(
@@ -886,6 +910,10 @@ def frc(self, other, cutoff=None, method="fft", plot=False):
             where `estimated_resolution` is in angstrom
             and FRC is a Numpy array of correlations.
         """
+        if not self._is_square:
+            raise NotImplementedError(
+                "`Image.frc` is not currently implemented for non-square images."
+            )
 
         if not isinstance(other, Image):
             raise TypeError(

src/aspire/source/micrograph.py

@@ -1,5 +1,6 @@
 import logging
 import os
+import warnings
 from abc import ABC, abstractmethod
 from glob import glob
 from pathlib import Path
@@ -20,7 +21,12 @@ class MicrographSource(ABC):
     def __init__(self, micrograph_count, micrograph_size, dtype, pixel_size=None):
         """ """
         self.micrograph_count = int(micrograph_count)
-        self.micrograph_size = int(micrograph_size)
+        # Expand single integer to 2-tuple
+        if isinstance(micrograph_size, int):
+            micrograph_size = (micrograph_size,) * 2
+        if len(micrograph_size) != 2:
+            raise ValueError("`micrograph_size` should be a integer or 2-tuple")
+        self.micrograph_size = tuple(micrograph_size)
         self.dtype = np.dtype(dtype)
         if pixel_size is not None:
             pixel_size = float(pixel_size)
@@ -34,7 +40,7 @@ def __repr__(self):
 
         :return: Returns a string description of instance.
         """
-        return f"{self.__class__.__name__} with {self.micrograph_count} {self.dtype.name} micrographs of size {self.micrograph_size}x{self.micrograph_size}"
+        return f"{self.__class__.__name__} with {self.micrograph_count} {self.dtype.name} micrographs of size {self.micrograph_size}"
 
     def __len__(self):
         """
@@ -142,14 +148,14 @@ def __init__(self, micrographs, dtype=None, pixel_size=None):
         if micrographs.ndim == 2:
             micrographs = micrographs[None, :, :]
 
-        if micrographs.ndim != 3 or (micrographs.shape[-2] != micrographs.shape[-1]):
+        if micrographs.ndim != 3:
             raise NotImplementedError(
-                f"Incompatible `micrographs` shape {micrographs.shape}, expects (count, L, L)"
+                f"Incompatible `micrographs` shape {micrographs.shape}, expects 2D or 3D array."
             )
 
         super().__init__(
             micrograph_count=micrographs.shape[0],
-            micrograph_size=micrographs.shape[-1],
+            micrograph_size=micrographs.shape[-2:],
             dtype=dtype or micrographs.dtype,
             pixel_size=pixel_size,
         )
@@ -201,15 +207,15 @@ def __init__(self, micrographs_path, dtype=None, pixel_size=None):
 
         # Load the first micrograph to infer shape/type
         # Size will be checked during on-the-fly loading of subsequent micrographs.
-        micrograph0 = Image.load(self.micrograph_files[0])
-        if micrograph0.pixel_size is not None and micrograph0.pixel_size != pixel_size:
-            raise ValueError(
-                f"Mismatched pixel size. {micrograph0.pixel_size} angstroms defined in {self.micrograph_files[0]}, but provided {pixel_size} angstroms."
-            )
+        micrograph0, _pixel_size = Image._load_raw(self.micrograph_files[0])
+        # Compare with user provided pixel size
+        if pixel_size is not None and _pixel_size != pixel_size:
+            msg = f"Mismatched pixel size. {_pixel_size} angstroms defined in {self.micrograph_files[0]}, but provided {pixel_size} angstroms."
+            warnings.warn(msg, UserWarning, stacklevel=2)
 
         super().__init__(
             micrograph_count=len(self.micrograph_files),
-            micrograph_size=micrograph0.resolution,
+            micrograph_size=micrograph0.shape[-2:],
             dtype=dtype or micrograph0.dtype,
             pixel_size=pixel_size,
         )
@@ -265,28 +271,27 @@ def _images(self, indices):
         # Initialize empty result
         n_micrographs = len(indices)
         micrographs = np.empty(
-            (n_micrographs, self.micrograph_size, self.micrograph_size),
+            (n_micrographs, *self.micrograph_size),
             dtype=self.dtype,
         )
         for i, ind in enumerate(indices):
             # Load the micrograph image from file
-            micrograph = Image.load(self.micrograph_files[ind])
+            micrograph, _pixel_size = Image._load_raw(self.micrograph_files[ind])
+
             # Assert size
-            if micrograph.resolution != self.micrograph_size:
+            if micrograph.shape != self.micrograph_size:
                 raise NotImplementedError(
                     f"Micrograph {ind} has inconsistent shape {micrograph.shape},"
-                    f" expected {(self.micrograph_size, self.micrograph_size)}."
+                    f" expected {self.micrograph_size}."
                 )
+
+            # Continually compare with initial pixel_size
+            if _pixel_size is not None and _pixel_size != self.pixel_size:
+                msg = f"Mismatched pixel size. {micrograph.pixel_size} angstroms defined in {self.micrograph_files[ind]}, but provided {self.pixel_size} angstroms."
+                warnings.warn(msg, UserWarning, stacklevel=2)
+
             # Assign to array, implicitly performs casting to dtype
-            micrographs[i] = micrograph.asnumpy()
-            # Assert pixel_size
-            if (
-                micrograph.pixel_size is not None
-                and micrograph.pixel_size != self.pixel_size
-            ):
-                raise ValueError(
-                    f"Mismatched pixel size. {micrograph.pixel_size} angstroms defined in {self.micrograph_files[ind]}, but provided {self.pixel_size} angstroms."
-                )
+            micrographs[i] = micrograph
 
         return Image(micrographs, pixel_size=self.pixel_size)
 
@@ -313,7 +318,7 @@ def __init__(
 
         :param volume: `Volume` instance to be used in `Simulation`.
              An `(L,L,L)` `Volume` will generate `(L,L)` particle images.
-        :param micrograph_size: Size of micrograph in pixels, defaults to 4096.
+        :param micrograph_size: Size of micrograph in pixels as integer or 2-tuple.  Defaults to 4096.
         :param micrograph_count: Number of micrographs to generate (integer). Defaults to 1.
         :param particles_per_micrograph: The amount of particles generated for each micrograph. Defaults to 10.
         :param particle_amplitudes: Optional, amplitudes to pass to `Simulation`.
@@ -354,7 +359,7 @@ def __init__(
 
         self.noise_adder = noise_adder
 
-        if self.particle_box_size > micrograph_size:
+        if self.particle_box_size > max(self.micrograph_size):
             raise ValueError(
                 "The micrograph size must be larger or equal to the `particle_box_size`."
             )
@@ -415,7 +420,7 @@ def __init__(
         else:
             if (
                 boundary < (-self.particle_box_size // 2)
-                or boundary > self.micrograph_size // 2
+                or boundary > max(self.micrograph_size) // 2
             ):
                 raise ValueError("Illegal boundary value.")
             self.boundary = boundary
@@ -505,8 +510,8 @@ def _set_mask(self):
         """
         self._mask = np.full(
             (
-                int(self.micrograph_size + 2 * self.pad),
-                int(self.micrograph_size + 2 * self.pad),
+                int(self.micrograph_size[0] + 2 * self.pad),
+                int(self.micrograph_size[1] + 2 * self.pad),
             ),
             False,
             dtype=bool,
@@ -547,7 +552,7 @@ def _clean_images(self, indices):
         # Initialize empty micrograph
         n_micrographs = len(indices)
         clean_micrograph = np.zeros(
-            (n_micrographs, self.micrograph_size, self.micrograph_size),
+            (n_micrographs, *self.micrograph_size),
             dtype=self.dtype,
         )
         # Pad the micrograph
@@ -579,8 +584,8 @@ def _clean_images(self, indices):
                 )
         clean_micrograph = clean_micrograph[
             :,
-            self.pad : self.micrograph_size + self.pad,
-            self.pad : self.micrograph_size + self.pad,
+            self.pad : self.micrograph_size[0] + self.pad,
+            self.pad : self.micrograph_size[1] + self.pad,
         ]
         return Image(clean_micrograph, pixel_size=self.pixel_size)
 

tests/test_array_image_source.py

@@ -81,7 +81,7 @@ def testArrayImageSourceNumpyError(self):
 
         # Test we raise with expected message from getter.
         with raises(RuntimeError, match=r"Creating Image object from Numpy.*"):
-            _ = ArrayImageSource(np.empty((3, 2, 1)))
+            _ = ArrayImageSource(np.empty((1)))
 
     def testArrayImageSourceAngGetterError(self):
         """

tests/test_image.py

@@ -88,9 +88,9 @@ def testRepr(get_mdim_images):
 
 
 def testNonSquare():
-    """Test that an irregular Image array raises."""
-    with raises(ValueError, match=r".* square .*"):
-        _ = Image(np.empty((4, 5)))
+    """Test that an irregular Image array does not raise."""
+    _ = Image(np.empty((4, 5)))
+    _ = Image(np.empty((3, 4, 5)))
 
 
 def testImShift(get_images, dtype):

tests/test_micrograph_simulation.py

@@ -85,11 +85,11 @@ def test_micrograph_source_has_correct_values(vol_fixture, micrograph_fixture):
     assert v.resolution == m.particle_box_size
     assert v == m.simulation.vols
     assert len(m) == m.micrograph_count
-    assert m.clean_images[0].shape[1] == m.micrograph_size
-    assert m.clean_images[0].shape[2] == m.micrograph_size
+    assert m.clean_images[0].shape[1] == m.micrograph_size[0]
+    assert m.clean_images[0].shape[2] == m.micrograph_size[1]
     assert (
         repr(m)
-        == f"{m.__class__.__name__} with {m.micrograph_count} {m.dtype.name} micrographs of size {m.micrograph_size}x{m.micrograph_size}"
+        == f"{m.__class__.__name__} with {m.micrograph_count} {m.dtype.name} micrographs of size {m.micrograph_size}"
     )
     _ = m.clean_images[:]
     _ = m.images[:]
@@ -136,9 +136,9 @@ def test_micrograph_centers_match(micrograph_fixture):
     for i, center in enumerate(centers):
         if (
             center[0] >= 0
-            and center[0] < m.micrograph_size
+            and center[0] < m.micrograph_size[0]
             and center[1] >= 0
-            and center[1] < m.micrograph_size
+            and center[1] < m.micrograph_size[1]
         ):
             assert m.clean_images[i // m.particles_per_micrograph].asnumpy()[0][
                 tuple(center)

tests/test_micrograph_source.py

@@ -271,12 +271,11 @@ def test_wrong_dim_micrograph_source():
 
 def test_rectangular_micrograph_source_array():
     """
-    Test non-square micrograph source raises.
+    Test non-square micrograph source does not raises.
     """
     # Test with Numpy array input
     imgs_np = np.empty((3, 7, 8))
-    with pytest.raises(RuntimeError, match=r"Incompatible.*"):
-        ArrayMicrographSource(imgs_np)
+    _ = ArrayMicrographSource(imgs_np)
 
 
 def test_rectangular_micrograph_source_files():