diff --git a/src/aspire/source/coordinates.py b/src/aspire/source/coordinates.py index 045ae5d8ae..7c8973792c 100644 --- a/src/aspire/source/coordinates.py +++ b/src/aspire/source/coordinates.py @@ -1,5 +1,6 @@ import logging import os +import warnings from abc import ABC, abstractmethod from collections import defaultdict from collections.abc import Iterable @@ -46,7 +47,19 @@ class CoordinateSource(ImageSource, ABC): This also allows the CoordinateSource to be saved to an `.mrcs` stack. """ - def __init__(self, files, particle_size, max_rows, B, symmetry_group): + def __init__( + self, files, particle_size, max_rows, B, symmetry_group, pixel_size=None + ): + """ + :param files: A list of tuples of the form (path_to_mrc, path_to_coord) + :param particle_size: Desired size of cropped particles (will override the size specified in coordinate file) + :param max_rows: Maximum number of particles to read. (If `None`, will attempt to load all particles) + :param B: CTF envelope decay factor + :param symmetry_group: A `SymmetryGroup` object or string corresponding to the symmetry of the molecule. + :param pixel_size: Pixel size of the images in angstroms. + Default `None` will attempt to infer `pixel_size` from + `CTFFilter` objects when available. + """ mrc_paths, coord_paths = [f[0] for f in files], [f[1] for f in files] # the particle_size parameter is the *user-specified* argument # and is used in self._populate_particles @@ -134,7 +147,14 @@ def __init__(self, files, particle_size, max_rows, B, symmetry_group): # total particles loaded (specific to this instance) logger.info(f"CoordinateSource object contains {n} particles.") - ImageSource.__init__(self, L=L, n=n, dtype=dtype, symmetry_group=symmetry_group) + ImageSource.__init__( + self, + L=L, + n=n, + dtype=dtype, + symmetry_group=symmetry_group, + pixel_size=pixel_size, + ) # map mrc indices to particle indices # i'th element contains a list of particle indices corresponding to i'th mrc @@ -385,6 +405,33 @@ def _extract_ctf(self, data_block): # convert defocus_ang from degrees to radians filter_params[:, 3] *= np.pi / 180.0 + # Check pixel_size + # Get pixel_sizes from CTFFilters + ctf_pixel_sizes = np.unique(filter_params[:, 6]) + # Compare with source.pixel_size if assigned + if (self.pixel_size is not None) and ( + not np.allclose(ctf_pixel_sizes, self.pixel_size) + ): + warnings.warn( + "Pixel size mismatch." + f"\n\tSource: {self.pixel_size}" + f"\n\tCTFs: {ctf_pixel_sizes}.", + stacklevel=2, + ) + # When source is not assigned we can try to assign it from CTF, + elif self.pixel_size is None: + # but only do this if all the CTFFilter pixel_sizes are consistent + if len(ctf_pixel_sizes) == 1: + self.pixel_size = ctf_pixel_sizes[0] # take the unique single element + logger.info( + f"Assigning source pixel_size={self.pixel_size} from CTFFilters." + ) + # otherwise let the user know + elif len(ctf_pixel_sizes) > 1: + logger.warning( + "Unable to assign source pixel_size from CTFFilters, multiple pixel_sizes found." + ) + # construct filters self.unique_filters = [ CTFFilter( @@ -516,16 +563,27 @@ def __init__( max_rows=None, B=0, symmetry_group=None, + pixel_size=None, ): """ :param files: A list of tuples of the form (path_to_mrc, path_to_coord) - :particle_size: Desired size of cropped particles (will override the size specified in coordinate file) + :param particle_size: Desired size of cropped particles (will override the size specified in coordinate file) :param max_rows: Maximum number of particles to read. (If `None`, will attempt to load all particles) + :param B: CTF envelope decay factor :param symmetry_group: A `SymmetryGroup` object or string corresponding to the symmetry of the molecule. + :param pixel_size: Pixel size of the images in angstroms. + Default `None` will attempt to infer `pixel_size` from + `CTFFilter` objects when available. """ # instantiate super CoordinateSource.__init__( - self, files, particle_size, max_rows, B, symmetry_group + self, + files, + particle_size, + max_rows, + B, + symmetry_group, + pixel_size=pixel_size, ) def _extract_box_size(self, box_file): @@ -629,17 +687,35 @@ class CentersCoordinateSource(CoordinateSource): Represents a data source consisting of micrographs and coordinate files specifying particle centers only. Files can be text (.coord) or STAR files. """ - def __init__(self, files, particle_size, max_rows=None, B=0, symmetry_group=None): + def __init__( + self, + files, + particle_size, + max_rows=None, + B=0, + symmetry_group=None, + pixel_size=None, + ): """ :param files: A list of tuples of the form (path_to_mrc, path_to_coord) - :particle_size: Desired size of cropped particles (mandatory) + :param particle_size: Desired size of cropped particles (mandatory) :param max_rows: Maximum number of particles to read. (If `None`, will attempt to load all particles) + :param B: CTF envelope decay factor :param symmetry_group: A `SymmetryGroup` object or string corresponding to the symmetry of the molecule. + :param pixel_size: Pixel size of the images in angstroms. + Default `None` will attempt to infer `pixel_size` from + `CTFFilter` objects when available. """ # instantiate super CoordinateSource.__init__( - self, files, particle_size, max_rows, B, symmetry_group + self, + files, + particle_size, + max_rows, + B, + symmetry_group, + pixel_size=pixel_size, ) def _validate_centers_file(self, coord_file): diff --git a/tests/test_coordinate_source.py b/tests/test_coordinate_source.py index 5304ef0b74..c8e8704e69 100644 --- a/tests/test_coordinate_source.py +++ b/tests/test_coordinate_source.py @@ -9,6 +9,7 @@ import mrcfile import numpy as np +import pytest from click.testing import CliRunner import tests.saved_test_data @@ -31,6 +32,7 @@ def setUp(self): self.original_mrc_path = str(test_path) # save test data root dir self.test_dir_root = os.path.dirname(self.original_mrc_path) + self.pixel_size = 1.23 # Used for generating and comparing metadata # We will construct a source with two micrographs and two coordinate # files by using the same micrograph, but dividing the coordinates @@ -136,6 +138,12 @@ def setUp(self): def tearDown(self): self.tmpdir.cleanup() + # This is a workaround to use a `pytest` fixture with `unittest` style cases. + # We use it below to capture and inspect the log + @pytest.fixture(autouse=True) + def inject_fixtures(self, caplog): + self._caplog = caplog + def createTestBoxFiles(self, centers, index): """ Create a .box file storing particle coordinates as @@ -250,6 +258,8 @@ def createFloatStarFile(self, centers): def createTestCtfFiles(self, index): """ Creates example ASPIRE-generated CTF files. + + Note two distinct pixel sizes. """ star_fp = os.path.join(self.data_folder, f"ctf{index+1}.star") # note that values are arbitrary and not representative of actual CTF data @@ -261,7 +271,7 @@ def createTestCtfFiles(self, index): "_rlnSphericalAberration": 700 + index, "_rlnAmplitudeContrast": 600 + index, "_rlnVoltage": 500 + index, - "_rlnMicrographPixelSize": 400 + index, + "_rlnMicrographPixelSize": self.pixel_size + index * 0.01, } blocks = OrderedDict({"root": params_dict}) starfile = StarFile(blocks=blocks) @@ -270,6 +280,8 @@ def createTestCtfFiles(self, index): def createTestRelionCtfFile(self, reverse_optics_block_rows=False): """ Creates example RELION-generated CTF file for a set of micrographs. + + Note uniform pixel size. """ star_fp = os.path.join(self.data_folder, "micrographs_ctf.star") blocks = OrderedDict() @@ -284,8 +296,8 @@ def createTestRelionCtfFile(self, reverse_optics_block_rows=False): ] # using same unique values as in createTestCtfFiles optics_block = [ - ["opticsGroup1", 1, 500.0, 700.0, 600.0, 400.0], - ["opticsGroup2", 2, 501.0, 701.0, 601.0, 401.0], + ["opticsGroup1", 1, 500.0, 700.0, 600.0, self.pixel_size], + ["opticsGroup2", 2, 501.0, 701.0, 601.0, self.pixel_size], ] # Since optics block rows are self-contained, # reversing their order should have no affect anywhere. @@ -530,8 +542,8 @@ def testWrongNumberCtfFiles(self): def testImportCtfFromList(self): src = BoxesCoordinateSource(self.files_box) src.import_aspire_ctf(self.ctf_files) - self._testCtfFilters(src) - self._testCtfMetadata(src) + self._testCtfFilters(src, uniform_pixel_sizes=False) + self._testCtfMetadata(src, uniform_pixel_sizes=False) def testImportCtfFromRelion(self): src = BoxesCoordinateSource(self.files_box) @@ -554,7 +566,7 @@ def testImportCtfFromRelionLegacy(self): self._testCtfFilters(src) self._testCtfMetadata(src) - def _testCtfFilters(self, src): + def _testCtfFilters(self, src, uniform_pixel_sizes=True): # there are two micrographs and two CTF files, so there should be two # unique CTF filters self.assertEqual(len(src.unique_filters), 2) @@ -565,7 +577,15 @@ def _testCtfFilters(self, src): self.assertTrue( np.allclose( np.array( - [1000.0, 900.0, 800.0 * np.pi / 180.0, 700.0, 600.0, 500.0, 400.0], + [ + 1000.0, + 900.0, + 800.0 * np.pi / 180.0, + 700.0, + 600.0, + 500.0, + self.pixel_size, + ], dtype=src.dtype, ), np.array( @@ -582,10 +602,21 @@ def _testCtfFilters(self, src): ) ) filter1 = src.unique_filters[1] + pixel_size1 = self.pixel_size + if not uniform_pixel_sizes: + pixel_size1 += 0.01 self.assertTrue( np.allclose( np.array( - [1001.0, 901.0, 801.0 * np.pi / 180.0, 701.0, 601.0, 501.0, 401.0], + [ + 1001.0, + 901.0, + 801.0 * np.pi / 180.0, + 701.0, + 601.0, + 501.0, + pixel_size1, + ], dtype=src.dtype, ), np.array( @@ -611,7 +642,7 @@ def _testCtfFilters(self, src): np.array_equal(np.where(src.filter_indices == 1)[0], np.arange(200, 400)) ) - def _testCtfMetadata(self, src): + def _testCtfMetadata(self, src, uniform_pixel_sizes=True): # ensure metadata is populated correctly when adding CTF info # __mrc_filepath mrc_fp_metadata = np.array( @@ -644,10 +675,13 @@ def _testCtfMetadata(self, src): ] ctf_metadata = np.zeros((src.n, len(ctf_cols)), dtype=src.dtype) ctf_metadata[:200] = np.array( - [1000.0, 900.0, 800.0 * np.pi / 180.0, 700.0, 600.0, 500.0, 400.0] + [1000.0, 900.0, 800.0 * np.pi / 180.0, 700.0, 600.0, 500.0, self.pixel_size] ) + pixel_size1 = self.pixel_size + if not uniform_pixel_sizes: + pixel_size1 += 0.01 ctf_metadata[200:400] = np.array( - [1001.0, 901.0, 801.0 * np.pi / 180.0, 701.0, 601.0, 501.0, 401.0] + [1001.0, 901.0, 801.0 * np.pi / 180.0, 701.0, 601.0, 501.0, pixel_size1] ) self.assertTrue(np.array_equal(ctf_metadata, src.get_metadata(ctf_cols))) @@ -700,6 +734,44 @@ def testCommand(self): self.assertTrue(result_star.exit_code == 0) self.assertTrue(result_preprocess.exit_code == 0) + def testPixelSizeWarning(self): + """ + Test source having a pixel size that conflicts with the CTFFilter instances. + """ + manual_pixel_size = 0.789 + src = BoxesCoordinateSource(self.files_box, pixel_size=manual_pixel_size) + # Capture and compare warning message + with pytest.warns(UserWarning, match=r".*Pixel size mismatch.*"): + src.import_relion_ctf(self.relion_ctf_file) + np.testing.assert_approx_equal(src.pixel_size, manual_pixel_size) + + def testMultiplePixelSizeWarning(self): + """ + Test source having multiple pixel sizes in CTFFilter instances. + """ + src = BoxesCoordinateSource(self.files_box) # pixel_size=None + # Capture and compare warning message + with self._caplog.at_level(logging.WARNING): + src.import_aspire_ctf(self.ctf_files) # not uniform_pixel_sizes + assert src.pixel_size is None + assert "multiple pixel_sizes found" in self._caplog.text + + def testPixelSize(self): + """ + Test explicitly providing correct pixel_size. + """ + src = BoxesCoordinateSource(self.files_box, pixel_size=self.pixel_size) + src.import_relion_ctf(self.relion_ctf_file) + np.testing.assert_approx_equal(src.pixel_size, self.pixel_size) + + def testPixelSizeNone(self): + """ + Test not providing pixel_size. + """ + src = BoxesCoordinateSource(self.files_box) + src.import_relion_ctf(self.relion_ctf_file) + np.testing.assert_approx_equal(src.pixel_size, self.pixel_size) + def create_test_rectangular_micrograph_and_star(tmp_path, voxel_size=(2.0, 2.0, 1.0)): # Create a rectangular micrograph (e.g., 128x256)