From df1b12dff15d11ff4e3e29596de5b5f52b8f20a0 Mon Sep 17 00:00:00 2001 From: elainehoml Date: Thu, 9 Mar 2023 16:19:04 +0000 Subject: [PATCH] specify aretomo path --- src/Ot2Rec/aretomo.py | 32 ++++++++++++++++++++++++++------ src/Ot2Rec/magicgui.py | 5 +++++ src/Ot2Rec/params.py | 1 + tests/test_aretomo_align.py | 36 ++++++++++++++++++++++++++++++++++++ tests/test_aretomo_recon.py | 33 +++++++++++++++++++++++++++++++++ 5 files changed, 101 insertions(+), 6 deletions(-) diff --git a/src/Ot2Rec/aretomo.py b/src/Ot2Rec/aretomo.py index 4184a77..760a191 100644 --- a/src/Ot2Rec/aretomo.py +++ b/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 diff --git a/src/Ot2Rec/magicgui.py b/src/Ot2Rec/magicgui.py index 73a7125..816c020 100644 --- a/src/Ot2Rec/magicgui.py +++ b/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() diff --git a/src/Ot2Rec/params.py b/src/Ot2Rec/params.py index 93f5862..5fd22ac 100644 --- a/src/Ot2Rec/params.py +++ b/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": { diff --git a/tests/test_aretomo_align.py b/tests/test_aretomo_align.py index fce18db..991f077 100644 --- a/tests/test_aretomo_align.py +++ b/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() + diff --git a/tests/test_aretomo_recon.py b/tests/test_aretomo_recon.py index ad06ff6..48ba80c 100644 --- a/tests/test_aretomo_recon.py +++ b/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()