diff --git a/src/pytom_tm/tmjob.py b/src/pytom_tm/tmjob.py index 595067da..8805885a 100644 --- a/src/pytom_tm/tmjob.py +++ b/src/pytom_tm/tmjob.py @@ -177,9 +177,17 @@ def __init__( self.ctf_data = ctf_data self.whiten_spectrum = 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(): + if self.whiten_spectrum: logging.info('Estimating whitening filter...') - weights = 1 / np.sqrt(power_spectrum_profile(read_mrc(self.tomogram))) + weights = 1 / np.sqrt( + power_spectrum_profile( + read_mrc(self.tomogram)[ + self.search_origin[0]: self.search_origin[0] + self.search_size[0], + self.search_origin[1]: self.search_origin[1] + self.search_size[1], + self.search_origin[2]: self.search_origin[2] + self.search_size[2] + ] + ) + ) weights /= weights.max() # scale to 1 np.save(self.whitening_filter, weights) diff --git a/tests/test_tmjob.py b/tests/test_tmjob.py index 0291c17f..8500630d 100644 --- a/tests/test_tmjob.py +++ b/tests/test_tmjob.py @@ -182,6 +182,18 @@ 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') + # load the whitening filter from previous job to compare against + whitening_filter = np.load(TEST_WHITENING_FILTER) + job = TMJob('0', 10, TEST_TOMOGRAM, TEST_TEMPLATE, TEST_MASK, TEST_DATA_DIR, + angle_increment='90.00', voxel_size=1., whiten_spectrum=True, search_y=[10, 90]) + new_whitening_filter = np.load(TEST_WHITENING_FILTER) + self.assertNotEqual(whitening_filter.shape, new_whitening_filter.shape, + msg='After reducing the search region along the largest dimension the whitening filter ' + 'should have less sampling points') + self.assertEqual(new_whitening_filter.shape, (max(job.search_size) // 2 + 1, ), + msg='The whitening filter does not have the expected size, it should be equal (x // 2) + 1, ' + 'where x is the largest dimension of the search box.') + # TMJob with none of these weighting options is tested in all other runs in this file. def test_custom_angular_search(self):