Allow custom angular sampling

src/bin/pytom_match_template.py

@@ -37,7 +37,10 @@ def main():
                              'CTF parameters are provided these will all be incorporated in the tilt-weighting.')
     parser.add_argument('--angular-search', type=str, required=True,
                         help='Options are: [7.00, 35.76, 19.95, 90.00, 18.00, '
-                             '12.85, 38.53, 11.00, 17.86, 25.25, 50.00, 3.00]')
+                             '12.85, 38.53, 11.00, 17.86, 25.25, 50.00, 3.00].\n'
+                             '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.')

src/pytom_tm/tmjob.py

@@ -151,7 +151,14 @@ def __init__(
             self.rotation_file = AVAILABLE_ROTATIONAL_SAMPLING[angle_increment][0]
             self.n_rotations = AVAILABLE_ROTATIONAL_SAMPLING[angle_increment][1]
         except KeyError:
-            raise TMJobError('Provided angular search is not available in  the default lists.')
+            possible_file_path = pathlib.Path(angle_increment)
+            if possible_file_path.exists() and possible_file_path.suffix == '.txt':
+                logging.info('Custom file provided for the angular search. Checking if it can be read...')
+                # load_angle_list will throw an error if it does not encounter three inputs per line
+                self.n_rotations = len(load_angle_list(possible_file_path))
+                self.rotation_file = possible_file_path
+            else:
+                raise TMJobError('Invalid angular search provided.')
 
         # missing wedge
         self.tilt_angles = tilt_angles

tests/test_tmjob.py

@@ -27,6 +27,7 @@
 TEST_MASK = TEST_DATA_DIR.joinpath('mask.mrc')
 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')
 
 
 class TestTMJob(unittest.TestCase):
@@ -92,6 +93,8 @@ def setUpClass(cls) -> None:
             job.voxel_size
         )
 
+        np.savetxt(TEST_CUSTOM_ANGULAR_SEARCH, np.random.rand(100, 3))
+
     @classmethod
     def tearDownClass(cls) -> None:
         TEST_MASK.unlink()
@@ -103,6 +106,7 @@ def tearDownClass(cls) -> None:
         TEST_TOMOGRAM.unlink()
         TEST_SCORES.unlink()
         TEST_ANGLES.unlink()
+        TEST_CUSTOM_ANGULAR_SEARCH.unlink()
         TEST_DATA_DIR.rmdir()
 
     def setUp(self):
@@ -145,6 +149,12 @@ def test_tm_job_copy(self):
             msg='Tomogram shape not correct in job, perhaps transpose issue?'
         )
 
+    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.)
+        self.assertEqual(job.rotation_file, TEST_CUSTOM_ANGULAR_SEARCH, msg='Passing a custom angular search file to '
+                                                                            'TMJob failed.')
+
     def test_tm_job_split_volume(self):
         sub_jobs = self.job.split_volume_search((1, 3, 1))
         for x in sub_jobs: