Merge pull request #90 from rosalindfranklininstitute/aretomo_specify…

…_exe

specify aretomo path

src/Ot2Rec/aretomo.py

@@ -122,10 +122,12 @@ def _get_aretomo_align_command(self, i):
             '0',
             '-OutBin',
             str(self.params['AreTomo_setup']['output_binning']),
-            '-DarkTol',
-            str(self.params['AreTomo_setup']['dark_tol']),
         ]
 
+        # Specify path to AreTomo if not using module loaded version
+        if len(self.params['System']['aretomo_path']) > 0:
+            cmd[0] = self.params['System']['aretomo_path']
+
         return cmd
 
     def _get_aretomo_recon_command(self, i):
@@ -158,6 +160,10 @@ def _get_aretomo_recon_command(self, i):
             cmd.append('-Wbp')
             cmd.append('1')
 
+        # Specify path to AreTomo if not using module loaded version
+        if len(self.params['System']['aretomo_path']) > 0:
+            cmd[0] = self.params['System']['aretomo_path']
+
         return cmd
 
     def _run_aretomo(self, i):
@@ -183,6 +189,11 @@ def _run_aretomo(self, i):
             }
             cmd.append('-OutImod')
             cmd.append(outimod_lookup[out_imod])
+
+        # Add darktol
+        if self.params['AreTomo_setup']['aretomo_mode'] != 1:
+            cmd.append("-DarkTol")
+            cmd.append(str(self.params['AreTomo_setup']['dark_tol']))
 
         # Add extra kwargs
         kwargs = self.params["AreTomo_kwargs"].keys()
@@ -415,14 +426,23 @@ def update_yaml(args):
 
         # Set output mrc
         output_lookup = {0: "_ali.mrc", 2: "_rec.mrc"}
+        # out_file_list = [
+        #     (f"{aretomo_params.params['System']['output_path']}/"
+        #      f"{os.path.splitext(os.path.basename(file))[0]}/"
+        #      f"{os.path.splitext(os.path.basename(file))[0]}"
+        #      f"{output_lookup[args['aretomo_mode']]}") for file in st_file_list
+        # ]
+
+        rootname = aretomo_params.params['System']['output_rootname']
+        suffix = aretomo_params.params['System']['output_suffix']
+        ext = output_lookup[args['aretomo_mode']]
         out_file_list = [
             (f"{aretomo_params.params['System']['output_path']}/"
-             f"{os.path.splitext(os.path.basename(file))[0]}/"
-             f"{os.path.splitext(os.path.basename(file))[0]}"
-             f"{output_lookup[args['aretomo_mode']]}") for file in st_file_list
+             f"{rootname}_{curr_ts:04d}{suffix}/"
+             f"{rootname}_{curr_ts:04d}{suffix}{ext}"
+            ) for curr_ts in ts_list
         ]
 
-
         aretomo_params.params["AreTomo_setup"]["output_mrc"] = out_file_list
 
     elif args["aretomo_mode"] == 1: # for reconstruction only

src/Ot2Rec/magicgui.py

@@ -820,6 +820,10 @@ def get_args_imod_route(
     dark_tol={
         "label": "Tolerance to remove dark images",
         "tooltip": "Default 0.7, low number = fewer images removed",
+    },
+    aretomo_path={
+        "label": "Path to AreTomo executable (leave blank if module loaded)",
+        "tooltip": "Ensure the path is to the correct version to match CUDA",
     }
 )
 def get_args_aretomo(
@@ -839,6 +843,7 @@ def get_args_aretomo(
         recon_algo="SART",
         out_imod="N/A",
         dark_tol=0.7,
+        aretomo_path="",
 ):
     return locals()
 

src/Ot2Rec/params.py

@@ -318,6 +318,7 @@ def new_aretomo_yaml(args):
             "output_path": str(args["output_path"]),
             "output_rootname": args["project_name"] if args["rootname"] == "" else args["rootname"],
             "output_suffix": args["suffix"],
+            "aretomo_path": args["aretomo_path"]
         },
 
         "AreTomo_setup": {

tests/test_aretomo_align.py

@@ -162,3 +162,39 @@ def test_aretomo_called(self, aretomo_mock):
         self.assertTrue(aretomo_mock.called)
 
         tmpdir.cleanup()
+
+    def test_aretomo_path_specified(self):
+        """ Tests that AreTomo path can be specified correctly in command """
+        # Create expected input
+        tmpdir = self._create_expected_folder_structure()
+        os.chdir(tmpdir.name)
+        args = self._create_expected_input_args()
+        self._create_expected_st_folder_structure(tmpdir)
+        args["aretomo_path"] = "/home/AreTomo_1.2.5_Cuda113_08-01-2022"
+
+        # Create yaml
+        aretomo.create_yaml(args)
+
+        # Read params
+        params = prmMod.read_yaml(
+            project_name="TS",
+            filename="./TS_aretomo_align.yaml",
+        )
+
+        # Run
+        logger = logMod.Logger("./o2r_aretomo_align.log")
+        aretomo_obj = aretomo.AreTomo(
+            project_name="TS",
+            params_in=params,
+            logger_in=logger
+        )        
+
+        cmd = aretomo_obj._get_aretomo_align_command(0)
+
+        self.assertEqual(
+            cmd[0],
+            "/home/AreTomo_1.2.5_Cuda113_08-01-2022"
+        )
+
+        tmpdir.cleanup()
+

tests/test_aretomo_recon.py

@@ -150,3 +150,36 @@ def test_aretomo_STA_output_folders_created(self, aretomo_mock):
         )
 
         tmpdir.cleanup()
+
+    def test_aretomo_path_specified(self):
+        # Create expected input
+        tmpdir = self._create_expected_folder_structure()
+        os.chdir(tmpdir.name)
+        args = self._create_expected_input_args()
+        args["aretomo_path"] = "/home/AreTomo_1.2.5_Cuda113_08-01-2022"
+
+        # Create yaml
+        aretomo.create_yaml(args)
+
+        # Read params
+        params = prmMod.read_yaml(
+            project_name="TS",
+            filename="./TS_aretomo_recon.yaml",
+        )
+
+        # Run
+        logger = logMod.Logger("./o2r_aretomo_recon.log")
+        aretomo_obj = aretomo.AreTomo(
+            project_name="TS",
+            params_in=params,
+            logger_in=logger
+        )
+
+        cmd = aretomo_obj._get_aretomo_recon_command(0)
+
+        self.assertEqual(
+            cmd[0],
+            "/home/AreTomo_1.2.5_Cuda113_08-01-2022"
+        )
+
+        tmpdir.cleanup()