add offset tests

Author: j-c-c

tests/test_relion_source.py

@@ -5,6 +5,7 @@
 import pytest
 
 from aspire.source import RelionSource
+from aspire.utils import RelionStarFile
 from aspire.volume import SymmetryGroup
 
 from .test_starfile_stack import StarFileTestCase
@@ -101,3 +102,70 @@ def test_pixel_size(caplog):
         src = RelionSource(starfile_im_pix_size, pixel_size=1.234)  # 1.4 in metadata
         # Ensure we prefer user provided
         np.testing.assert_allclose(src.pixel_size, 1.234)
+
+
+def test_offsets_conversion():
+    """
+    Check that offset convention gets converted to Relion >= 3.1 convention.
+    """
+    starfile = os.path.join(DATA_DIR, "sample_particles_relion30.star")
+
+    # Extract pixel valued offsets from starfile prior to source instantiation.
+    metadata = RelionStarFile(starfile).get_merged_data_block()
+    pixel_offsets = np.column_stack((metadata["_rlnOriginX"], metadata["_rlnOriginY"]))
+
+    # Create Relion Source and extract offsets from metadata using updated field names.
+    src = RelionSource(starfile)
+    angst_offsets = src.get_metadata(["_rlnOriginXAngst", "_rlnOriginYAngst"])
+
+    # Check that old convention offset fields have been removed.
+    assert "_rlnOriginX" not in src._metadata
+    assert "_rlnOriginY" not in src._metadata
+
+    # Check that offsets in metadata match up to pixel/angstrom conversion.
+    np.testing.assert_allclose(angst_offsets / src.pixel_size, pixel_offsets)
+
+    # src.offsets should still return pixel valued offsets.
+    np.testing.assert_allclose(src.offsets, pixel_offsets)
+
+
+def test_offsets():
+    """
+    Check that offsets are loaded properly with starfile field _rlnOriginX(Y)Angst.
+    """
+    # This starfile has offsets stored with angstrom values as _rlnOriginX(Y)Angst.
+    starfile = os.path.join(DATA_DIR, "sample_particles_relion31.star")
+
+    # Create a RelionSource
+    src = RelionSource(starfile)
+
+    # Check offsets are angstrom valued in metadata and correspond to src.offsets.
+    angst_offsets = src.get_metadata(["_rlnOriginXAngst", "_rlnOriginYAngst"])
+    np.testing.assert_allclose(src.offsets * src.pixel_size, angst_offsets)
+
+
+def test_offsets_save(tmp_path):
+    """
+    Test that saving a RelionSource that was loaded with pixel offsets
+    saves with angstrom valued offsets.
+    """
+    # Starfile with pixel offsets.
+    starfile = os.path.join(DATA_DIR, "sample_particles_relion30.star")
+
+    # Extract pixel valued offsets from starfile prior to source instantiation.
+    metadata = RelionStarFile(starfile).get_merged_data_block()
+    pixel_offsets = np.column_stack((metadata["_rlnOriginX"], metadata["_rlnOriginY"]))
+
+    # Create and RelionSource and save to starfile.
+    src = RelionSource(starfile)
+    save_path = tmp_path / "test_file.star"
+    src.save(save_path)
+
+    # Saved starfile should have angstrom valued offsets.
+    metadata = RelionStarFile(save_path).get_merged_data_block()
+    angst_offsets = np.column_stack(
+        (metadata["_rlnOriginXAngst"], metadata["_rlnOriginYAngst"])
+    )
+
+    # Check saved offsets match original up to pixel_size scaling.
+    np.testing.assert_allclose(angst_offsets / src.pixel_size, pixel_offsets)