Fix CoordinateSource (allow rectangular micrographs)

src/aspire/image/image.py

@@ -598,6 +598,10 @@ def save(self, mrcs_filepath, overwrite=None):
         if self.stack_ndim > 1:
             raise NotImplementedError("`save` is currently limited to 1D image stacks.")
 
+        data = self._data.astype(np.float32)
+        if self.n_images == 1:
+            data = data[0]
+
         if overwrite is None and os.path.exists(mrcs_filepath):
             # If the file exists, append a timestamp to the old file and rename it
             _ = rename_with_timestamp(mrcs_filepath)
@@ -606,13 +610,13 @@ def save(self, mrcs_filepath, overwrite=None):
 
         with mrcfile.new(mrcs_filepath, overwrite=overwrite) as mrc:
             # original input format (the image index first)
-            mrc.set_data(self._data.astype(np.float32))
+            mrc.set_data(data)
             # Note assigning voxel_size must come after `set_data`
             if self.pixel_size is not None:
                 mrc.voxel_size = self.pixel_size
 
     @staticmethod
-    def load(filepath, dtype=None):
+    def _load_raw(filepath, dtype=None):
         """
         Load raw data from supported files.
 
@@ -621,7 +625,9 @@ def load(filepath, dtype=None):
         :param filepath: File path (string).
         :param dtype: Optionally force cast to `dtype`.
              Default dtype is inferred from the file contents.
-        :return: numpy array of image data.
+        :returns:
+            - numpy array of image data.
+            - pixel size
         """
 
         # Get the file extension
@@ -640,6 +646,23 @@ def load(filepath, dtype=None):
         if dtype is not None:
             im = im.astype(dtype, copy=False)
 
+        return im, pixel_size
+
+    @staticmethod
+    def load(filepath, dtype=None):
+        """
+        Load raw data from supported files.
+
+        Currently MRC and TIFF are supported.
+
+        :param filepath: File path (string).
+        :param dtype: Optionally force cast to `dtype`.
+             Default dtype is inferred from the file contents.
+        :return: Image instance
+        """
+        # Load raw data from filepath with pixel size
+        im, pixel_size = Image._load_raw(filepath, dtype=dtype)
+
         # Return as Image instance
         return Image(im, pixel_size=pixel_size)
 
@@ -867,7 +890,7 @@ def _vx_array_to_size(vx):
         #   checks uniformity.
         if isinstance(vx, np.recarray):
             if vx.x != vx.y:
-                raise ValueError(f"Voxel sizes are not uniform: {vx}")
+                logger.warning(f"Voxel sizes are not uniform: {vx}")
             vx = vx.x
 
         # Convert `0` to `None`

src/aspire/source/coordinates.py

@@ -308,7 +308,7 @@ def _get_mrc_shapes(self):
 
         mrc_shapes = np.zeros((self.num_micrographs, 2), dtype=int)
         for i, mrc in enumerate(self.mrc_paths):
-            mrc_shapes[i, :] = Image.load(mrc).resolution
+            mrc_shapes[i, :] = Image._load_raw(mrc)[0].shape
 
         return mrc_shapes
 
@@ -469,7 +469,7 @@ def _images(self, indices):
         # their origin micrograph
         for mrc_index, coord_list in grouped.items():
             # Load file as 2D numpy array.
-            arr = Image.load(self.mrc_paths[mrc_index]).asnumpy()[0]
+            arr = Image._load_raw(self.mrc_paths[mrc_index])[0]
 
             # create iterable of the coordinates in this mrc
             # we don't need to worry about exhausting this iter

tests/test_coordinate_source.py

@@ -1,3 +1,4 @@
+import logging
 import os
 import random
 import shutil
@@ -17,6 +18,8 @@
 from aspire.storage import StarFile
 from aspire.utils import RelionStarFile, importlib_path
 
+logger = logging.getLogger(__name__)
+
 
 class CoordinateSourceTestCase(TestCase):
     def setUp(self):
@@ -696,3 +699,50 @@ def testCommand(self):
         self.assertTrue(result_coord.exit_code == 0)
         self.assertTrue(result_star.exit_code == 0)
         self.assertTrue(result_preprocess.exit_code == 0)
+
+
+def create_test_rectangular_micrograph_and_star(tmp_path, voxel_size=(2.0, 2.0, 1.0)):
+    # Create a rectangular micrograph (e.g., 128x256)
+    data = np.random.rand(128, 256).astype(np.float32)
+    mrc_path = tmp_path / "test_micrograph.mrc"
+
+    with mrcfile.new(mrc_path, overwrite=True) as mrc:
+        mrc.set_data(data)
+        mrc.voxel_size = voxel_size
+
+    # Two sample coordinates
+    coordinates = [(50.0, 30.0), (200.0, 100.0)]
+
+    # Write a simple STAR file
+    star_path = tmp_path / "test_coordinates.star"
+    with open(star_path, "w") as f:
+        f.write("data_particles\n\n")
+        f.write("loop_\n")
+        f.write("_rlnCoordinateX #1\n")
+        f.write("_rlnCoordinateY #2\n")
+        for x, y in coordinates:
+            f.write(f"{x:.1f} {y:.1f}\n")
+
+    # Pack files into a list of tuples for consumption by CoordinatSource
+    file_list = [(mrc_path, star_path)]
+
+    return file_list
+
+
+def test_rectangular_coordinate_source(tmp_path):
+    file_list = create_test_rectangular_micrograph_and_star(tmp_path)
+
+    # Check we can instantiate a CoordinateSource with a rectangular micrograph.
+    coord_src = CentersCoordinateSource(file_list, particle_size=32)
+
+    # Check we can access images.
+    _ = coord_src.images[:]
+
+
+def test_coordinate_source_pixel_warning(tmp_path, caplog):
+    # Create micrograph with mismatched pixel dimensions.
+    vx = (2.3, 2.1, 1.0)
+    file_list = create_test_rectangular_micrograph_and_star(tmp_path, voxel_size=vx)
+    with caplog.at_level(logging.WARNING):
+        _ = CentersCoordinateSource(file_list, particle_size=32)
+        assert "Voxel sizes are not uniform" in caplog.text