Repository: ComputationalCryoEM/ASPIRE-Python

Issue:
OrientedSource single mode `save`
The single save mode in OrientedSource appears to not work  I took a peek at the code, it looked like the logic has a bug in it.  Need to add a test of this feature, then patch the code.

In a pipeline, I have the following:
```
...
avgs.save(class_avg_fn, save_mode="single", overwrite=True)
....
oriented_src.save(oriented_fn, save_mode="single", overwrite=True)
....
```

Which yields the files below.  Can see the first set of files output as expected; the second set should be a single star and single mrcs.

```
...
10073_preprocessed_129px.star
10073_preprocessed_129px_0_138898.mrcs
...
10073_legacy_cls_avgs_m50_129px.star
10073_legacy_cls_avgs_m50_129px_0_138898.mrcs
10073_legacy_oriented_class_averages_129px.star
10073_legacy_oriented_class_averages_129px_0_511.mrcs
10073_legacy_oriented_class_averages_129px_512_1023.mrcs
10073_legacy_oriented_class_averages_129px_1024_1535.mrcs
10073_legacy_oriented_class_averages_129px_1536_2047.mrcs
10073_legacy_oriented_class_averages_129px_2048_2559.mrcs
10073_legacy_oriented_class_averages_129px_2560_3019.mrcs
...
```

Hints from issue comments:

Solution patch:
diff --git a/src/aspire/source/image.py b/src/aspire/source/image.py
index f529e153..d084deb2 100644
--- a/src/aspire/source/image.py
+++ b/src/aspire/source/image.py
@@ -1695,7 +1695,7 @@ class OrientedSource(IndexedSource):
     def save_metadata(self, starfile_filepath, batch_size=512, save_mode=None):
         self._orient()
         return super().save_metadata(
-            starfile_filepath, batch_size=batch_size, save_mode=None
+            starfile_filepath, batch_size=batch_size, save_mode=save_mode
         )
 
     def get_metadata(


Test patch:
diff --git a/tests/test_oriented_source.py b/tests/test_oriented_source.py
index bdc0b3f0..ebabb818 100644
--- a/tests/test_oriented_source.py
+++ b/tests/test_oriented_source.py
@@ -1,10 +1,13 @@
+import glob
 import logging
+import os
+import tempfile
 
 import numpy as np
 import pytest
 
 from aspire.abinitio import CLSymmetryC3C4, CLSymmetryCn, CLSync3N, CLSyncVoting
-from aspire.source import OrientedSource, Simulation
+from aspire.source import OrientedSource, RelionSource, Simulation
 from aspire.volume import CnSymmetricVolume
 
 logger = logging.getLogger(__name__)
@@ -16,6 +19,7 @@ ESTIMATOR_SYMMETRY = [
     (CLSymmetryC3C4, "C4"),
     pytest.param((CLSymmetryCn, "C6"), marks=pytest.mark.expensive),
 ]
+SAVE_MODES = [None, "single"]
 
 
 def src_fixture_id(params):
@@ -46,6 +50,11 @@ def src_fixture(request):
     return og_src, oriented_src
 
 
+@pytest.fixture(params=SAVE_MODES, ids=lambda x: f"save_mode={x}")
+def save_mode(request):
+    return request.param
+
+
 def test_repr(src_fixture):
     og_src, oriented_src = src_fixture
 
@@ -122,3 +131,42 @@ def test_lazy_evaluation(src_fixture, caplog):
     assert msg not in caplog.text
     _ = oriented_src.rotations
     assert msg in caplog.text
+
+
+def test_save(src_fixture, save_mode):
+    """
+    Test save function and save_mode.
+    """
+
+    src = src_fixture[1]
+
+    # Make a fresh tmp_dir
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        # Construct file path
+        fn = os.path.join(tmp_dir, f"test_oriented_source-save_mode-{save_mode}.star")
+
+        # Sanity check test configuration
+        assert src.n % 2 == 0, "Test expects even number of source images"
+
+        # Configure save
+        #   Ensure we have a (small) unit test sized `batch_size`
+        batch_size = src.n // 2
+        # Set expected output file counts
+        if save_mode == "single":
+            num_batches = 1
+        else:
+            num_batches = src.n // batch_size
+
+        # Save
+        src.save(fn, save_mode=save_mode, batch_size=batch_size)
+
+        # Load from saved `fn`
+        reloaded_src = RelionSource(fn)
+
+        # Assert reloaded image data identical
+        np.testing.assert_allclose(reloaded_src.images[:], src.images[:])
+
+        # Assert the correct number of mrcs files were created
+        mrc_filelist = glob.glob(os.path.join(tmp_dir, "*.mrcs"))
+        num_mrc_files = len(mrc_filelist)
+        assert num_mrc_files == num_batches, "Incorrect number of mrcs files"