diff --git a/setup.py b/setup.py index 94e15cd4..142100b2 100644 --- a/setup.py +++ b/setup.py @@ -8,7 +8,7 @@ name='pytom-template-matching-gpu', packages=['pytom_tm', 'pytom_tm.angle_lists'], package_dir={'': 'src'}, - version='0.3.0', # for versioning definition see https://semver.org/ + version='0.3.1', # for versioning definition see https://semver.org/ description='GPU template matching from PyTOM as a lightweight pip package', long_description=long_description, long_description_content_type='text/markdown', @@ -24,7 +24,8 @@ 'tqdm', 'mrcfile', 'starfile', - 'importlib_resources' + 'importlib_resources', + 'packaging', ], extras_require={ 'plotting': ['matplotlib', 'seaborn'] diff --git a/src/bin/pytom_match_template.py b/src/bin/pytom_match_template.py index 26dcb7fe..e240f9c7 100644 --- a/src/bin/pytom_match_template.py +++ b/src/bin/pytom_match_template.py @@ -41,9 +41,10 @@ def main(): 'Alternatively, a .txt file can be provided with three Euler angles (in radians) per ' 'line that define the angular search. Angle format is ZXZ anti-clockwise (see: ' 'https://www.ccpem.ac.uk/user_help/rotation_conventions.php).') - parser.add_argument('--rotational-symmetry', type=int, required=False, action=LargerThanZero, default=1, - help='Integer value indicating the rotational symmetry of the template. The length of the ' - 'rotation search will be shortened through division by this value.') + parser.add_argument('--z-axis-rotational-symmetry', type=int, required=False, action=LargerThanZero, default=1, + help='Integer value indicating the rotational symmetry of the template around the z-axis. The ' + 'length of the rotation search will be shortened through division by this value. Only ' + 'works for template symmetry around the z-axis.') parser.add_argument('-s', '--volume-split', nargs=3, type=int, required=False, default=[1, 1, 1], help='Split the volume into smaller parts for the search, can be relevant if the volume does ' 'not fit into GPU memory. Format is x y z, e.g. --volume-split 1 2 1') @@ -125,7 +126,7 @@ def main(): dose_accumulation=args.dose_accumulation, ctf_data=ctf_params, whiten_spectrum=args.spectral_whitening, - rotational_symmetry=args.rotational_symmetry, + rotational_symmetry=args.z_axis_rotational_symmetry, ) score_volume, angle_volume = run_job_parallel(job, tuple(args.volume_split), args.gpu_ids) diff --git a/src/pytom_tm/__init__.py b/src/pytom_tm/__init__.py index 2b5a53b7..22d3ec9e 100644 --- a/src/pytom_tm/__init__.py +++ b/src/pytom_tm/__init__.py @@ -1,3 +1,6 @@ +from importlib import metadata +__version__ = metadata.version('pytom-template-matching-gpu') + try: import cupy except (ModuleNotFoundError, ImportError): diff --git a/src/pytom_tm/angles.py b/src/pytom_tm/angles.py index d5d9be4f..ca9e15bb 100644 --- a/src/pytom_tm/angles.py +++ b/src/pytom_tm/angles.py @@ -22,14 +22,15 @@ v[0] = ANGLE_LIST_DIR.joinpath(v[0]) -def load_angle_list(file_name: pathlib.Path) -> list[tuple[float, float, float]]: +def load_angle_list(file_name: pathlib.Path, sort_angles: bool = True) -> list[tuple[float, float, float]]: with open(str(file_name)) as fstream: lines = fstream.readlines() angle_list = [tuple(map(float, x.strip().split(' '))) for x in lines] if not all([len(a) == 3 for a in angle_list]): raise ValueError('Invalid angle file provided, each line should have 3 ZXZ Euler angles!') - else: - return angle_list + if sort_angles: + angle_list.sort() # angle list needs to be sorted otherwise symmetry reduction cannot be used! + return angle_list def convert_euler( diff --git a/src/pytom_tm/extract.py b/src/pytom_tm/extract.py index 3437efa6..c85c6e3b 100644 --- a/src/pytom_tm/extract.py +++ b/src/pytom_tm/extract.py @@ -1,3 +1,4 @@ +from packaging import version import pandas as pd import numpy as np import numpy.typing as npt @@ -36,7 +37,10 @@ def extract_particles( score_volume = read_mrc(job.output_dir.joinpath(f'{job.tomo_id}_scores.mrc')) angle_volume = read_mrc(job.output_dir.joinpath(f'{job.tomo_id}_angles.mrc')) - angle_list = load_angle_list(job.rotation_file) + angle_list = load_angle_list( + job.rotation_file, + sort_angles=version.parse(job.pytom_tm_version_number) > version.parse('0.3.0') + ) # mask edges of score volume score_volume[0: particle_radius_px, :, :] = -1 @@ -61,7 +65,7 @@ def extract_particles( search_space = ( # wherever the score volume has not been explicitly set to -1 is the size of the search region (score_volume > -1).sum() * - job.n_rotations + int(np.ceil(job.n_rotations / job.rotational_symmetry)) ) cut_off = erfcinv((2 * n_false_positives) / search_space) * np.sqrt(2) * sigma logging.info(f'cut off for particle extraction: {cut_off}') diff --git a/src/pytom_tm/tmjob.py b/src/pytom_tm/tmjob.py index 6cd56ca9..95ef1d03 100644 --- a/src/pytom_tm/tmjob.py +++ b/src/pytom_tm/tmjob.py @@ -1,4 +1,6 @@ from __future__ import annotations +from importlib import metadata +from packaging import version import pathlib import copy import numpy as np @@ -39,6 +41,8 @@ def load_json_to_tmjob(file_name: pathlib.Path) -> TMJob: ctf_data=data.get('ctf_data', None), whiten_spectrum=data.get('whiten_spectrum', False), rotational_symmetry=data.get('rotational_symmetry', 1), + # if version number is not in the .json, it must be 0.3.0 or older + pytom_tm_version_number=data.get('pytom_tm_version_number', '0.3.0'), ) job.rotation_file = pathlib.Path(data['rotation_file']) job.whole_start = data['whole_start'] @@ -79,7 +83,8 @@ def __init__( dose_accumulation: Optional[list[float, ...]] = None, ctf_data: Optional[list[dict, ...]] = None, whiten_spectrum: bool = False, - rotational_symmetry: int = 1 + rotational_symmetry: int = 1, + pytom_tm_version_number: str = metadata.version('pytom-template-matching-gpu') ): self.mask = mask self.mask_is_spherical = mask_is_spherical @@ -171,12 +176,12 @@ def __init__( self.dose_accumulation = dose_accumulation self.ctf_data = ctf_data self.whiten_spectrum = whiten_spectrum - - if self.whiten_spectrum: + self.whitening_filter = self.output_dir.joinpath(f'{self.tomo_id}_whitening_filter.npy') + if self.whiten_spectrum and not self.whitening_filter.exists(): logging.info('Estimating whitening filter...') weights = 1 / np.sqrt(power_spectrum_profile(read_mrc(self.tomogram))) weights /= weights.max() # scale to 1 - np.save(self.output_dir.joinpath('whitening_filter.npy'), weights) + np.save(self.whitening_filter, weights) # Job details self.job_key = job_key @@ -188,6 +193,9 @@ def __init__( self.log_level = log_level + # version number of the job + self.pytom_tm_version_number = pytom_tm_version_number + def copy(self) -> TMJob: return copy.deepcopy(self) @@ -395,7 +403,7 @@ def start_job( self.low_pass, self.high_pass ) * (profile_to_weighting( - np.load(self.output_dir.joinpath('whitening_filter.npy')), + np.load(self.whitening_filter), search_volume.shape ) if self.whiten_spectrum else 1)).astype(np.float32) @@ -415,7 +423,7 @@ def start_job( accumulated_dose_per_tilt=self.dose_accumulation, ctf_params_per_tilt=self.ctf_data ) * (profile_to_weighting( - np.load(self.output_dir.joinpath('whitening_filter.npy')), + np.load(self.whitening_filter), self.template_shape ) if self.whiten_spectrum else 1)).astype(np.float32) @@ -447,7 +455,10 @@ def start_job( int(np.ceil(self.n_rotations / self.rotational_symmetry)), self.steps_slice )) - angle_list = load_angle_list(self.rotation_file)[slice( + angle_list = load_angle_list( + self.rotation_file, + sort_angles=version.parse(self.pytom_tm_version_number) > version.parse('0.3.0') + )[slice( self.start_slice, int(np.ceil(self.n_rotations / self.rotational_symmetry)), self.steps_slice diff --git a/tests/test_template_matching.py b/tests/test_template_matching.py index 62befd5c..802822ff 100644 --- a/tests/test_template_matching.py +++ b/tests/test_template_matching.py @@ -40,7 +40,7 @@ def test_search(self): self.assertEqual(angle_id, angle_volume[ind]) self.assertSequenceEqual(loc, ind) self.assertEqual(stats['search_space'], 256000000, msg='Search space should exactly equal this value') - self.assertAlmostEqual(stats['std'], 0.005175, places=6, + self.assertAlmostEqual(stats['std'], 0.005175, places=5, msg='Standard deviation of the search should be almost equal') diff --git a/tests/test_tmjob.py b/tests/test_tmjob.py index 3f0a9a19..5ea84732 100644 --- a/tests/test_tmjob.py +++ b/tests/test_tmjob.py @@ -172,7 +172,7 @@ def test_tm_job_split_volume(self): score, angle = self.job.merge_sub_jobs() ind = np.unravel_index(score.argmax(), score.shape) - self.assertTrue(score.max() > 0.934, msg='lcc max value lower than expected') + self.assertTrue(score.max() > 0.931, msg='lcc max value lower than expected') self.assertEqual(ANGLE_ID, angle[ind]) self.assertSequenceEqual(LOCATION, ind) @@ -204,7 +204,7 @@ def test_tm_job_split_angles(self): score, angle = self.job.merge_sub_jobs() ind = np.unravel_index(score.argmax(), score.shape) - self.assertTrue(score.max() > 0.934, msg='lcc max value lower than expected') + self.assertTrue(score.max() > 0.931, msg='lcc max value lower than expected') self.assertEqual(ANGLE_ID, angle[ind]) self.assertSequenceEqual(LOCATION, ind) @@ -217,7 +217,7 @@ def test_parallel_manager(self): score, angle = run_job_parallel(self.job, volume_splits=(1, 3, 1), gpu_ids=[0]) ind = np.unravel_index(score.argmax(), score.shape) - self.assertTrue(score.max() > 0.934, msg='lcc max value lower than expected') + self.assertTrue(score.max() > 0.931, msg='lcc max value lower than expected') self.assertEqual(ANGLE_ID, angle[ind]) self.assertSequenceEqual(LOCATION, ind)