From 07443d48b607505f4aa7e8d4d5708158b1b0b81f Mon Sep 17 00:00:00 2001 From: McHaillet Date: Thu, 9 Nov 2023 13:42:46 +0100 Subject: [PATCH 1/6] remove writing of debug volumes as it interfers with testing; reorder the filter generation in the start of the template matching job to make it respond properly to input -- users should be able to apply a bandpass filter without providing tilt angles --- src/pytom_tm/tmjob.py | 82 +++++++++++++++++++++---------------------- 1 file changed, 40 insertions(+), 42 deletions(-) diff --git a/src/pytom_tm/tmjob.py b/src/pytom_tm/tmjob.py index 95ef1d03..2589251f 100644 --- a/src/pytom_tm/tmjob.py +++ b/src/pytom_tm/tmjob.py @@ -393,61 +393,59 @@ def start_job( read_mrc(self.mask) ) - # create weighting, include missing wedge and band-pass filters - template_wedge = None - if self.tilt_angles is not None: - # convolute tomo with whitening filter and bandpass - tomo_filter = (create_gaussian_band_pass( - search_volume.shape, - self.voxel_size, + # init tomogram and template weighting + tomo_filter, template_wedge = 1, 1 + # first generate bandpass filters + if not (self.low_pass is None and self.high_pass is None): + tomo_filter *= create_gaussian_band_pass( + search_volume.shape, + self.voxel_size, + self.low_pass, + self.high_pass + ).astype(np.float32) + template_wedge *= create_gaussian_band_pass( + self.template_shape, + self.voxel_size, self.low_pass, self.high_pass - ) * (profile_to_weighting( + ).astype(np.float32) + + # then multiply with optional whitening filters + if self.whiten_spectrum: + tomo_filter *= profile_to_weighting( np.load(self.whitening_filter), search_volume.shape - ) if self.whiten_spectrum else 1)).astype(np.float32) - - # we always apply a binary wedge (with optional band pass) to the volume to remove empty regions - search_volume = np.real(irfftn(rfftn(search_volume) * tomo_filter, s=search_volume.shape)) + ).astype(np.float32) + template_wedge *= profile_to_weighting( + np.load(self.whitening_filter), + self.template_shape + ).astype(np.float32) - # get template wedge - template_wedge = (create_wedge( + # if tilt angles are provided we can create wedge filters + if self.tilt_angles is not None: + # for the tomogram a binary wedge is generated to explicitly set the missing wedge region to 0 + tomo_filter *= create_wedge( + search_volume.shape, + self.tilt_angles, + self.voxel_size, + cut_off_radius=1., + angles_in_degrees=True, + tilt_weighting=False + ) + # for the template a binary or per-tilt-weighted wedge is generated depending on the options + template_wedge *= create_wedge( self.template_shape, self.tilt_angles, self.voxel_size, cut_off_radius=1., angles_in_degrees=True, - low_pass=self.low_pass, - high_pass=self.high_pass, tilt_weighting=self.tilt_weighting, accumulated_dose_per_tilt=self.dose_accumulation, ctf_params_per_tilt=self.ctf_data - ) * (profile_to_weighting( - np.load(self.whitening_filter), - self.template_shape - ) if self.whiten_spectrum else 1)).astype(np.float32) - - if logging.root.level == logging.DEBUG: - write_mrc( - self.output_dir.joinpath('debug_tomo_filter.mrc'), - tomo_filter, - voxel_size=self.voxel_size - ) - write_mrc( - self.output_dir.joinpath('debug_search_volume.mrc'), - search_volume, - voxel_size=self.voxel_size - ) - write_mrc( - self.output_dir.joinpath('debug_template_convoluted.mrc'), - np.fft.irfftn(np.fft.rfftn(template) * template_wedge, s=template.shape).astype(np.float32), - voxel_size=self.voxel_size - ) - write_mrc( - self.output_dir.joinpath('debug_3d_ctf.mrc'), - template_wedge, - voxel_size=self.voxel_size - ) + ).astype(np.float32) + + # apply the optional band pass and whitening filter to the search region + search_volume = np.real(irfftn(rfftn(search_volume) * tomo_filter, s=search_volume.shape)) # load rotation search angle_ids = list(range( From 39dbae69230673b918c63250e7beed7372b6b036 Mon Sep 17 00:00:00 2001 From: McHaillet Date: Thu, 9 Nov 2023 13:51:04 +0100 Subject: [PATCH 2/6] add weight test that checks tilt weighting without defocus and dose --- tests/test_weights.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/tests/test_weights.py b/tests/test_weights.py index 5a9f650f..ca17ea66 100644 --- a/tests/test_weights.py +++ b/tests/test_weights.py @@ -224,6 +224,7 @@ def test_create_wedge(self): self.assertEqual(structured_wedge.dtype, np.float32, msg='Wedge with band-pass does not have expected dtype') + # test shapes of wedges weights = create_wedge(self.volume_shape_even, self.tilt_angles, self.voxel_size * 3, tilt_weighting=True, low_pass=40, accumulated_dose_per_tilt=ACCUMULATED_DOSE, @@ -237,6 +238,14 @@ def test_create_wedge(self): self.assertEqual(weights.shape, self.reduced_uneven_shape_3d, msg='3D CTF does not have the correct reduced fourier shape.') + # test parameter flexibility of tilt_weighted wedge + weights = create_wedge(self.volume_shape_even, self.tilt_angles, self.voxel_size * 3, + tilt_weighting=True, low_pass=self.low_pass, + accumulated_dose_per_tilt=None, + ctf_params_per_tilt=None) + self.assertEqual(weights.shape, self.reduced_even_shape_3d, + msg='Tilt weighted wedge should also work without defocus and dose info.') + def test_ctf(self): ctf_raw = create_ctf( self.volume_shape_even, From 317d852eb786b590de31fc1b6cd7138842b071c4 Mon Sep 17 00:00:00 2001 From: McHaillet Date: Thu, 9 Nov 2023 14:28:07 +0100 Subject: [PATCH 3/6] added tests for the tmjob with all the different weighting options and combintations --- tests/test_tmjob.py | 30 ++++++++++++++++++++++++++++++ tests/test_weights.py | 20 ++++++++++---------- 2 files changed, 40 insertions(+), 10 deletions(-) diff --git a/tests/test_tmjob.py b/tests/test_tmjob.py index 5ea84732..aeab20e3 100644 --- a/tests/test_tmjob.py +++ b/tests/test_tmjob.py @@ -10,6 +10,7 @@ from pytom_tm.tmjob import TMJob, TMJobError from pytom_tm.io import read_mrc, write_mrc, UnequalSpacingError from pytom_tm.extract import extract_particles +from test_weights import CTF_PARAMS, ACCUMULATED_DOSE, TILT_ANGLES TOMO_SHAPE = (100, 107, 59) @@ -149,6 +150,35 @@ def test_tm_job_copy(self): msg='Tomogram shape not correct in job, perhaps transpose issue?' ) + def test_tm_job_weighting_options(self): + # run with all options + job = TMJob('0', 10, TEST_TOMOGRAM, TEST_TEMPLATE, TEST_MASK, TEST_DATA_DIR, + angle_increment='90.00', voxel_size=1., low_pass=10, high_pass=100, + dose_accumulation=ACCUMULATED_DOSE, ctf_data=CTF_PARAMS, tilt_angles=TILT_ANGLES, + whiten_spectrum=True, tilt_weighting=True) + score, angle = job.start_job(0, return_volumes=True) + self.assertEqual(score.shape, job.tomo_shape, msg='TMJob with all options failed') + + # run with only tilt weighting (in test_weights different options are tested for create_wedge) + job = TMJob('0', 10, TEST_TOMOGRAM, TEST_TEMPLATE, TEST_MASK, TEST_DATA_DIR, + angle_increment='90.00', voxel_size=1., + dose_accumulation=ACCUMULATED_DOSE, ctf_data=CTF_PARAMS, tilt_angles=TILT_ANGLES, + tilt_weighting=True) + score, angle = job.start_job(0, return_volumes=True) + self.assertEqual(score.shape, job.tomo_shape, msg='TMJob with only wedge creation failed') + + # run with only bandpass (in test_weights bandpass option are tested) + job = TMJob('0', 10, TEST_TOMOGRAM, TEST_TEMPLATE, TEST_MASK, TEST_DATA_DIR, + angle_increment='90.00', voxel_size=1., low_pass=10, high_pass=100) + score, angle = job.start_job(0, return_volumes=True) + self.assertEqual(score.shape, job.tomo_shape, msg='TMJob with only band-pass failed') + + # run with only spectrum whitening (in test_weights the whitening filter is tested + job = TMJob('0', 10, TEST_TOMOGRAM, TEST_TEMPLATE, TEST_MASK, TEST_DATA_DIR, + angle_increment='90.00', voxel_size=1., whiten_spectrum=True) + score, angle = job.start_job(0, return_volumes=True) + self.assertEqual(score.shape, job.tomo_shape, msg='TMJob with only whitening filter failed') + def test_custom_angular_search(self): job = TMJob('0', 10, TEST_TOMOGRAM, TEST_TEMPLATE, TEST_MASK, TEST_DATA_DIR, angle_increment=TEST_CUSTOM_ANGULAR_SEARCH, voxel_size=1.) diff --git a/tests/test_weights.py b/tests/test_weights.py index ca17ea66..82e34fea 100644 --- a/tests/test_weights.py +++ b/tests/test_weights.py @@ -91,6 +91,7 @@ 65.895 ''' ACCUMULATED_DOSE = [float(x.strip()) for x in DOSE_FILE.split('\n') if x != ''] +TILT_ANGLES = list(range(-51, 54, 3)) class TestWeights(unittest.TestCase): @@ -98,7 +99,6 @@ def setUp(self): self.volume_shape_even = (10, 10, 10) self.volume_shape_uneven = (11, 11, 11) self.volume_shape_irregular = (7, 12, 6) - self.tilt_angles = list(range(-51, 54, 3)) self.voxel_size = 3.34 self.low_pass = 10 self.high_pass = 50 @@ -180,23 +180,23 @@ def test_create_wedge(self): 'equal to 0'): create_wedge( self.volume_shape_even, - self.tilt_angles, + TILT_ANGLES, 0. ) with self.assertRaises(ValueError, msg='Create wedge should raise ValueError if cut_off_radius is smaller or ' 'equal to 0'): create_wedge( self.volume_shape_even, - self.tilt_angles, + TILT_ANGLES, 1., cut_off_radius=0. ) # create test wedges - structured_wedge = create_wedge(self.volume_shape_even, self.tilt_angles, 1., tilt_weighting=True) - symmetric_wedge = create_wedge(self.volume_shape_even, [self.tilt_angles[0], self.tilt_angles[-1]], + structured_wedge = create_wedge(self.volume_shape_even, TILT_ANGLES, 1., tilt_weighting=True) + symmetric_wedge = create_wedge(self.volume_shape_even, [TILT_ANGLES[0], TILT_ANGLES[-1]], 1., tilt_weighting=False) - asymmetric_wedge = create_wedge(self.volume_shape_even, [self.tilt_angles[0], self.tilt_angles[-2]], + asymmetric_wedge = create_wedge(self.volume_shape_even, [TILT_ANGLES[0], TILT_ANGLES[-2]], 1., tilt_weighting=False) self.assertEqual(structured_wedge.shape, self.reduced_even_shape_3d, @@ -217,7 +217,7 @@ def test_create_wedge(self): self.assertTrue(np.sum((symmetric_wedge != asymmetric_wedge) * 1) != 0, msg='Symmetric and asymmetric wedge should be different!') - structured_wedge = create_wedge(self.volume_shape_even, self.tilt_angles, self.voxel_size, tilt_weighting=True, + structured_wedge = create_wedge(self.volume_shape_even, TILT_ANGLES, self.voxel_size, tilt_weighting=True, cut_off_radius=1., low_pass=self.low_pass, high_pass=self.high_pass) self.assertEqual(structured_wedge.shape, self.reduced_even_shape_3d, msg='Wedge with band-pass does not have expected output shape') @@ -225,13 +225,13 @@ def test_create_wedge(self): msg='Wedge with band-pass does not have expected dtype') # test shapes of wedges - weights = create_wedge(self.volume_shape_even, self.tilt_angles, self.voxel_size * 3, + weights = create_wedge(self.volume_shape_even, TILT_ANGLES, self.voxel_size * 3, tilt_weighting=True, low_pass=40, accumulated_dose_per_tilt=ACCUMULATED_DOSE, ctf_params_per_tilt=CTF_PARAMS) self.assertEqual(weights.shape, self.reduced_even_shape_3d, msg='3D CTF does not have the correct reduced fourier shape.') - weights = create_wedge(self.volume_shape_uneven, self.tilt_angles, self.voxel_size * 3, + weights = create_wedge(self.volume_shape_uneven, TILT_ANGLES, self.voxel_size * 3, tilt_weighting=True, low_pass=40, accumulated_dose_per_tilt=ACCUMULATED_DOSE, ctf_params_per_tilt=CTF_PARAMS) @@ -239,7 +239,7 @@ def test_create_wedge(self): msg='3D CTF does not have the correct reduced fourier shape.') # test parameter flexibility of tilt_weighted wedge - weights = create_wedge(self.volume_shape_even, self.tilt_angles, self.voxel_size * 3, + weights = create_wedge(self.volume_shape_even, TILT_ANGLES, self.voxel_size * 3, tilt_weighting=True, low_pass=self.low_pass, accumulated_dose_per_tilt=None, ctf_params_per_tilt=None) From 1ef1b24fda4b8d834d137380c89c3bc466c9b873 Mon Sep 17 00:00:00 2001 From: McHaillet Date: Thu, 9 Nov 2023 14:29:10 +0100 Subject: [PATCH 4/6] add informative comment about testing the job without weighting --- tests/test_tmjob.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/test_tmjob.py b/tests/test_tmjob.py index aeab20e3..8ffebb34 100644 --- a/tests/test_tmjob.py +++ b/tests/test_tmjob.py @@ -179,6 +179,8 @@ def test_tm_job_weighting_options(self): score, angle = job.start_job(0, return_volumes=True) self.assertEqual(score.shape, job.tomo_shape, msg='TMJob with only whitening filter failed') + # TMJob with none of these weighting options is tested in all other runs in this file. + def test_custom_angular_search(self): job = TMJob('0', 10, TEST_TOMOGRAM, TEST_TEMPLATE, TEST_MASK, TEST_DATA_DIR, angle_increment=TEST_CUSTOM_ANGULAR_SEARCH, voxel_size=1.) From ba6801770ccb1b9049d3d26714c67e1e68e4a2df Mon Sep 17 00:00:00 2001 From: McHaillet Date: Thu, 9 Nov 2023 14:37:22 +0100 Subject: [PATCH 5/6] optional clean up of whitening filter --- tests/test_tmjob.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/test_tmjob.py b/tests/test_tmjob.py index 8ffebb34..6f574503 100644 --- a/tests/test_tmjob.py +++ b/tests/test_tmjob.py @@ -29,6 +29,7 @@ TEST_SCORES = TEST_DATA_DIR.joinpath('tomogram_scores.mrc') TEST_ANGLES = TEST_DATA_DIR.joinpath('tomogram_angles.mrc') TEST_CUSTOM_ANGULAR_SEARCH = TEST_DATA_DIR.joinpath('custom_angular_search.txt') +TEST_WHITENING_FILTER = TEST_DATA_DIR.joinpath('tomogram_whitening_filter.npy') class TestTMJob(unittest.TestCase): @@ -108,6 +109,9 @@ def tearDownClass(cls) -> None: TEST_SCORES.unlink() TEST_ANGLES.unlink() TEST_CUSTOM_ANGULAR_SEARCH.unlink() + # the whitening filter might not exist if the job with spectrum whitening failed, so the unlinking needs to + # allow this (with missing_ok=True) to ensure clean up of the test directory + TEST_WHITENING_FILTER.unlink(missing_ok=True) TEST_DATA_DIR.rmdir() def setUp(self): From 90bf79024206afe578453f3a04a2812cfae64448 Mon Sep 17 00:00:00 2001 From: Marten <58044494+McHaillet@users.noreply.github.com> Date: Fri, 10 Nov 2023 13:50:20 +0100 Subject: [PATCH 6/6] forgot to set tomo wedge type to float32 Co-authored-by: Sander Roet --- src/pytom_tm/tmjob.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pytom_tm/tmjob.py b/src/pytom_tm/tmjob.py index 2589251f..26a5cfb6 100644 --- a/src/pytom_tm/tmjob.py +++ b/src/pytom_tm/tmjob.py @@ -431,7 +431,7 @@ def start_job( cut_off_radius=1., angles_in_degrees=True, tilt_weighting=False - ) + ).astype(np.float32) # for the template a binary or per-tilt-weighted wedge is generated depending on the options template_wedge *= create_wedge( self.template_shape,